Skip to content

Commit c8392a9

Browse files
patelchaitanypatelchaitany
committed
Phase 1: Added basic open ai competible search api in the vector store
Signed-off-by: Chaitany patel <patelchaitany93@gmail.com>
1 parent e9dc0e0 commit c8392a9

13 files changed

Lines changed: 721 additions & 35 deletions

File tree

‎sdk/python/feast/feature_server.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,11 @@
5050
from feast.data_source import PushMode
5151
from feast.errors import (
5252
FeastError,
53+
FeatureViewNotFoundException,
5354
)
5455
from feast.feast_object import FeastObject
5556
from feast.feature_view_utils import get_feature_view_from_feature_store
57+
from feast.filter_models import ComparisonFilter, CompoundFilter
5658
from feast.permissions.action import WRITE, AuthzedAction
5759
from feast.permissions.security_manager import assert_permissions
5860
from feast.permissions.server.rest import inject_user_details
@@ -109,6 +111,40 @@ class GetOnlineDocumentsRequest(BaseModel):
109111
query: Optional[List[float]] = None
110112
query_string: Optional[str] = None
111113
api_version: Optional[int] = 1
114+
filters: Optional[Union[ComparisonFilter, CompoundFilter]] = None
115+
116+
117+
class OpenAISearchMetadata(BaseModel):
118+
features_to_retrieve: Optional[List[str]] = None
119+
content_field: Optional[str] = None
120+
121+
122+
class OpenAIComparisonFilter(BaseModel):
123+
key: str
124+
type: str
125+
value: Union[str, int, float, bool, List[Union[str, int]]]
126+
127+
128+
class OpenAICompoundFilter(BaseModel):
129+
type: str
130+
filters: List[Union[OpenAIComparisonFilter, "OpenAICompoundFilter"]]
131+
132+
133+
OpenAICompoundFilter.model_rebuild()
134+
135+
136+
class OpenAIRankingOptions(BaseModel):
137+
ranker: Optional[str] = None
138+
score_threshold: Optional[float] = None
139+
140+
141+
class OpenAISearchRequest(BaseModel):
142+
query: Union[str, List[str]]
143+
filters: Optional[Union[OpenAIComparisonFilter, OpenAICompoundFilter]] = None
144+
max_num_results: Optional[int] = 10
145+
ranking_options: Optional[OpenAIRankingOptions] = None
146+
rewrite_query: Optional[bool] = None
147+
metadata: Optional[OpenAISearchMetadata] = None
112148

113149

114150
class ChatMessage(BaseModel):
@@ -390,6 +426,8 @@ async def retrieve_online_documents(
390426
)
391427
if request.api_version == 2 and request.query_string is not None:
392428
read_params["query_string"] = request.query_string
429+
if request.api_version == 2 and request.filters is not None:
430+
read_params["filters"] = request.filters
393431

394432
if request.api_version == 2:
395433
response = await run_in_threadpool(
@@ -408,6 +446,51 @@ async def retrieve_online_documents(
408446
)
409447
return ORJSONResponse(content=response_dict)
410448

449+
@app.post(
450+
"/v1/vector_stores/{vector_store_id}/search",
451+
dependencies=[Depends(inject_user_details)],
452+
)
453+
async def openai_vector_store_search(
454+
vector_store_id: str,
455+
request: OpenAISearchRequest,
456+
) -> ORJSONResponse:
457+
with feast_metrics.track_request_latency(
458+
"/v1/vector_stores/{vector_store_id}/search"
459+
):
460+
try:
461+
result = await run_in_threadpool(
462+
lambda: store.retrieve_online_documents_openai(
463+
vector_store_id=vector_store_id,
464+
query=request.query,
465+
max_num_results=request.max_num_results or 10,
466+
filters=(
467+
request.filters.model_dump() if request.filters else None
468+
),
469+
ranking_options=(
470+
request.ranking_options.model_dump()
471+
if request.ranking_options
472+
else None
473+
),
474+
rewrite_query=request.rewrite_query,
475+
features_to_retrieve=(
476+
request.metadata.features_to_retrieve
477+
if request.metadata
478+
else None
479+
),
480+
)
481+
)
482+
except FeatureViewNotFoundException:
483+
return ORJSONResponse(
484+
status_code=404,
485+
content={
486+
"error": {
487+
"message": f"No vector store found with id '{vector_store_id}'",
488+
"type": "not_found_error",
489+
}
490+
},
491+
)
492+
return ORJSONResponse(content=result)
493+
411494
@app.post("/push", dependencies=[Depends(inject_user_details)])
412495
async def push(request: PushFeaturesRequest) -> Response:
413496
with feast_metrics.track_request_latency("/push"):

‎sdk/python/feast/feature_store.py‎

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@
6969
from feast.feast_object import FeastObject
7070
from feast.feature_service import FeatureService
7171
from feast.feature_view import DUMMY_ENTITY, DUMMY_ENTITY_NAME, FeatureView
72+
from feast.filter_models import ComparisonFilter, CompoundFilter, convert_dict_to_filter
7273
from feast.inference import (
7374
update_data_sources_with_inferred_event_timestamp_col,
7475
update_feature_views_with_inferred_features_and_entities,
@@ -2677,6 +2678,7 @@ def retrieve_online_documents_v2(
26772678
text_weight: float = 0.5,
26782679
image_weight: float = 0.5,
26792680
combine_strategy: str = "weighted_sum",
2681+
filters: Optional[Union[ComparisonFilter, CompoundFilter]] = None,
26802682
) -> OnlineResponse:
26812683
"""
26822684
Retrieves the top k closest document features. Note, embeddings are a subset of features.
@@ -2829,8 +2831,164 @@ def retrieve_online_documents_v2(
28292831
top_k,
28302832
distance_metric,
28312833
query_string,
2834+
filters,
28322835
)
28332836

2837+
def retrieve_online_documents_openai(
2838+
self,
2839+
vector_store_id: str,
2840+
query: Union[str, List[str]],
2841+
max_num_results: int = 10,
2842+
filters: Optional[Dict[str, Any]] = None,
2843+
ranking_options: Optional[Dict[str, Any]] = None,
2844+
rewrite_query: Optional[bool] = None,
2845+
features_to_retrieve: Optional[List[str]] = None,
2846+
) -> Dict[str, Any]:
2847+
"""
2848+
OpenAI-compatible vector store search.
2849+
2850+
Accepts a raw query string, optionally embeds it via LiteLLM
2851+
(when ``query_embedding_model`` is configured in feature_store.yaml),
2852+
and returns results in OpenAI's ``vector_store.search_results.page``
2853+
format.
2854+
2855+
Args:
2856+
vector_store_id: Feature view name (maps to the OpenAI
2857+
``vector_store_id`` path parameter).
2858+
query: Natural language query string, or list of strings.
2859+
max_num_results: Maximum number of results to return.
2860+
filters: OpenAI-compatible filters (accepted but not yet
2861+
applied).
2862+
ranking_options: OpenAI-compatible ranking options (accepted
2863+
but not yet applied).
2864+
rewrite_query: Whether to rewrite the query (accepted but
2865+
not yet applied).
2866+
features_to_retrieve: Specific feature names to return.
2867+
If None, all features from the feature view are used.
2868+
2869+
Returns:
2870+
Dict matching the OpenAI ``vector_store.search_results.page``
2871+
schema.
2872+
2873+
Examples:
2874+
Keyword search (no embedding model configured)::
2875+
2876+
result = store.retrieve_online_documents_openai(
2877+
vector_store_id="city_embeddings",
2878+
query="cities in California",
2879+
max_num_results=5,
2880+
)
2881+
2882+
Vector search (embedding model configured in YAML)::
2883+
2884+
# feature_store.yaml has:
2885+
# feature_server:
2886+
# query_embedding_model: text-embedding-3-small
2887+
result = store.retrieve_online_documents_openai(
2888+
vector_store_id="product_embeddings",
2889+
query="wireless audio device",
2890+
max_num_results=3,
2891+
features_to_retrieve=["name", "description"],
2892+
)
2893+
"""
2894+
feature_view = self.get_feature_view(vector_store_id)
2895+
2896+
if features_to_retrieve:
2897+
feature_names = features_to_retrieve
2898+
else:
2899+
feature_names = [f.name for f in feature_view.features]
2900+
2901+
features = [f"{feature_view.name}:{name}" for name in feature_names]
2902+
query_text = query if isinstance(query, str) else " ".join(query)
2903+
2904+
embed_cfg = self.config.embedding_model
2905+
if embed_cfg is None:
2906+
raise ValueError(
2907+
"embedding_model is not configured in feature_store.yaml. "
2908+
"Add an 'embedding_model' section with at least a 'model' "
2909+
"field to use retrieve_online_documents_openai.\n"
2910+
"Example:\n"
2911+
" embedding_model:\n"
2912+
" model: text-embedding-3-small\n"
2913+
" api_key: sk-..."
2914+
)
2915+
2916+
try:
2917+
from litellm import embedding as litellm_embedding
2918+
except ImportError:
2919+
raise ImportError(
2920+
"litellm is required for query embedding. "
2921+
"Install with: pip install litellm"
2922+
)
2923+
2924+
litellm_kwargs: Dict[str, Any] = {
2925+
"model": embed_cfg.model,
2926+
"input": [query_text],
2927+
}
2928+
if embed_cfg.api_key:
2929+
litellm_kwargs["api_key"] = embed_cfg.api_key
2930+
if embed_cfg.api_base:
2931+
litellm_kwargs["api_base"] = embed_cfg.api_base
2932+
if embed_cfg.api_version:
2933+
litellm_kwargs["api_version"] = embed_cfg.api_version
2934+
if embed_cfg.dimensions:
2935+
litellm_kwargs["dimensions"] = embed_cfg.dimensions
2936+
2937+
embed_response = litellm_embedding(**litellm_kwargs)
2938+
query_embedding = embed_response.data[0]["embedding"]
2939+
2940+
typed_filters: Optional[Union[ComparisonFilter, CompoundFilter]] = None
2941+
if filters is not None:
2942+
typed_filters = convert_dict_to_filter(filters)
2943+
2944+
response = self.retrieve_online_documents_v2(
2945+
features=features,
2946+
query=query_embedding,
2947+
top_k=max_num_results,
2948+
filters=typed_filters,
2949+
)
2950+
2951+
response_dict = response.to_dict()
2952+
2953+
result_data = []
2954+
if response_dict:
2955+
first_key = next(iter(response_dict))
2956+
num_rows = len(response_dict.get(first_key, []))
2957+
for i in range(num_rows):
2958+
score = 0.0
2959+
attributes: Dict[str, Any] = {}
2960+
content_parts: List[Dict[str, str]] = []
2961+
2962+
for key, values in response_dict.items():
2963+
val = values[i] if i < len(values) else None
2964+
if key == "distance":
2965+
score = float(val) if val is not None else 0.0
2966+
else:
2967+
attributes[key] = val
2968+
if isinstance(val, str):
2969+
content_parts.append({"type": "text", "text": val})
2970+
2971+
result_data.append(
2972+
{
2973+
"file_id": f"{vector_store_id}_{i}",
2974+
"filename": vector_store_id,
2975+
"score": score,
2976+
"attributes": attributes,
2977+
"content": content_parts
2978+
if content_parts
2979+
else [{"type": "text", "text": str(attributes)}],
2980+
}
2981+
)
2982+
2983+
search_query = query if isinstance(query, list) else [query]
2984+
return {
2985+
"object": "vector_store.search_results.page",
2986+
"search_query": search_query,
2987+
"data": result_data,
2988+
"has_more": False,
2989+
"next_page": None,
2990+
}
2991+
28342992
def _retrieve_from_online_store(
28352993
self,
28362994
provider: Provider,
@@ -2893,6 +3051,7 @@ def _retrieve_from_online_store_v2(
28933051
top_k: int,
28943052
distance_metric: Optional[str],
28953053
query_string: Optional[str],
3054+
filters: Optional[Union[ComparisonFilter, CompoundFilter]] = None,
28963055
) -> OnlineResponse:
28973056
"""
28983057
Search and return document features from the online document store.
@@ -2909,6 +3068,7 @@ def _retrieve_from_online_store_v2(
29093068
top_k=top_k,
29103069
distance_metric=distance_metric,
29113070
query_string=query_string,
3071+
filters=filters,
29123072
)
29133073

29143074
entity_key_dict: Dict[str, List[ValueProto]] = {}

‎sdk/python/feast/filter_models.py‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
from typing import Any, Dict, List, Literal, Optional, Union
2+
3+
from pydantic import BaseModel
4+
5+
6+
class ComparisonFilter(BaseModel):
7+
"""A filter that compares a metadata field against a value.
8+
9+
:param type: The comparison operator to apply
10+
:param key: The metadata field name to filter on
11+
:param value: The value to compare against
12+
"""
13+
14+
type: Literal["eq", "ne", "gt", "gte", "lt", "lte", "in", "nin"]
15+
key: str
16+
value: Any
17+
18+
19+
class CompoundFilter(BaseModel):
20+
"""A filter that combines multiple filters with a logical operator.
21+
22+
:param type: The logical operator ("and" requires all filters match,
23+
"or" requires any filter matches)
24+
:param filters: The list of filters to combine
25+
"""
26+
27+
type: Literal["and", "or"]
28+
filters: List[Union[ComparisonFilter, "CompoundFilter"]]
29+
30+
31+
CompoundFilter.model_rebuild()
32+
33+
FilterType = Optional[Union[ComparisonFilter, CompoundFilter]]
34+
35+
36+
def convert_dict_to_filter(
37+
filter_dict: Dict[str, Any],
38+
) -> Union[ComparisonFilter, CompoundFilter]:
39+
"""Convert a raw dict (e.g. from OpenAI-compatible JSON) into a typed filter object."""
40+
filter_type = filter_dict.get("type")
41+
if filter_type in ("and", "or"):
42+
return CompoundFilter(
43+
type=filter_type,
44+
filters=[convert_dict_to_filter(f) for f in filter_dict["filters"]],
45+
)
46+
return ComparisonFilter(
47+
type=filter_dict["type"],
48+
key=filter_dict["key"],
49+
value=filter_dict["value"],
50+
)

0 commit comments

Comments
 (0)