From f5686306b1392ff6815596e7a66cfad3b838b6cd Mon Sep 17 00:00:00 2001 From: Vanshika Vanshika Date: Thu, 16 Apr 2026 20:11:04 +0530 Subject: [PATCH] feat(cli): add RAG template as opt-in option for feast init MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a new RAG (Retrieval-Augmented Generation) template that can be selected via `feast init -t rag`. The template provides a City Q&A demo using Feast for feature management and Milvus for vector search. The default `feast init` behavior is unchanged — it continues to use the local template, preserving operator compatibility and avoiding heavy dependencies (pymilvus, torch, transformers) in the default getting-started flow. Fixes #5264 Signed-off-by: Vanshika Vanshika rh-pre-commit.version: 2.3.2 rh-pre-commit.check-secrets: ENABLED --- sdk/python/feast/cli/cli.py | 1 + sdk/python/feast/templates/rag/README.md | 99 +++++++++++ sdk/python/feast/templates/rag/__init__.py | 1 + sdk/python/feast/templates/rag/bootstrap.py | 56 +++++++ .../templates/rag/feature_repo/__init__.py | 1 + .../rag/feature_repo/example_repo.py | 137 ++++++++++++++++ .../rag/feature_repo/feature_store.yaml | 20 +++ .../rag/feature_repo/test_workflow.py | 155 ++++++++++++++++++ sdk/python/feast/templates/rag/gitignore | 33 ++++ 9 files changed, 503 insertions(+) create mode 100644 sdk/python/feast/templates/rag/README.md create mode 100644 sdk/python/feast/templates/rag/__init__.py create mode 100644 sdk/python/feast/templates/rag/bootstrap.py create mode 100644 sdk/python/feast/templates/rag/feature_repo/__init__.py create mode 100644 sdk/python/feast/templates/rag/feature_repo/example_repo.py create mode 100644 sdk/python/feast/templates/rag/feature_repo/feature_store.yaml create mode 100644 sdk/python/feast/templates/rag/feature_repo/test_workflow.py create mode 100644 sdk/python/feast/templates/rag/gitignore diff --git a/sdk/python/feast/cli/cli.py b/sdk/python/feast/cli/cli.py index 1e461af4a28..114082e1385 100644 --- a/sdk/python/feast/cli/cli.py +++ b/sdk/python/feast/cli/cli.py @@ -475,6 +475,7 @@ def materialize_incremental_command( "milvus", "ray", "ray_rag", + "rag", "pytorch_nlp", ], case_sensitive=False, diff --git a/sdk/python/feast/templates/rag/README.md b/sdk/python/feast/templates/rag/README.md new file mode 100644 index 00000000000..a6aa153d6a2 --- /dev/null +++ b/sdk/python/feast/templates/rag/README.md @@ -0,0 +1,99 @@ +# City Information Q&A — RAG Demo with Feast + +A complete Retrieval-Augmented Generation (RAG) demo using Feast for feature management and Milvus for vector search. + + +## Project Structure + +``` +rag/ +├── feature_repo/ +│ ├── data/ +│ │ └── city_wikipedia_summaries_with_embeddings.parquet # Sample data (US cities) +│ ├── example_repo.py # Entity, Feature Views, Feature Service definitions +│ ├── feature_store.yaml # Feast config (Milvus online store, file offline store) +│ └── test_workflow.py # End-to-end demo: apply → materialize → search +└── README.md +``` + +## Quick Start + +### 1. Initialize the template + +```bash +feast init -t rag my_city_qa +cd my_city_qa/feature_repo +``` + +### 2. Install dependencies + +```bash +pip install feast torch transformers pymilvus +``` + +### 3. Apply feature definitions + +```bash +feast apply +``` + +### 4. Explore in the Feast UI + +```bash +feast ui +``` + +### 5. Run the demo workflow + +```bash +python test_workflow.py +``` + + +## Key Commands + +| Command | Description | +|---------|-------------| +| `feast apply` | Register entities, feature views, and feature services | +| `feast materialize --disable-event-timestamp` | Load parquet data into the online store (Milvus) for vector search. Optionally add `-v city_summary_embeddings -v city_metadata` to materialize only those views. | +| `feast feature-views list` | List registered feature views | +| `feast entities list` | List registered entities | +| `feast feature-services list` | List registered feature services | +| `feast ui` | Start the Feast UI at http://localhost:8888 | + + +## Architecture + +``` +┌─────────────────────────────────────────────────────────────────┐ +│ City Q&A Pipeline │ +├─────────────────────────────────────────────────────────────────┤ +│ │ +│ User Question │ +│ │ │ +│ ▼ │ +│ ┌─────────────┐ │ +│ │ Embed Query │ (MiniLM 384-dim) │ +│ └─────────────┘ │ +│ │ │ +│ ▼ │ +│ ┌─────────────────────────────────────┐ │ +│ │ city_summary_embeddings (Milvus) │ ← Vector Search │ +│ │ - vector (COSINE similarity) │ │ +│ │ - sentence_chunks │ │ +│ └─────────────────────────────────────┘ │ +│ │ │ +│ ▼ (top-k city_ids) │ +│ ┌─────────────────────────────────────┐ │ +│ │ city_metadata (Feast Online Store) │ ← Metadata Lookup │ +│ │ - state │ │ +│ │ - wiki_summary │ │ +│ └─────────────────────────────────────┘ │ +│ │ │ +│ ▼ │ +│ ┌─────────────┐ │ +│ │ LLM Answer │ (optional: GPT/Claude) │ +│ └─────────────┘ │ +│ │ +└─────────────────────────────────────────────────────────────────┘ +``` diff --git a/sdk/python/feast/templates/rag/__init__.py b/sdk/python/feast/templates/rag/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/sdk/python/feast/templates/rag/__init__.py @@ -0,0 +1 @@ + diff --git a/sdk/python/feast/templates/rag/bootstrap.py b/sdk/python/feast/templates/rag/bootstrap.py new file mode 100644 index 00000000000..a84fb23a639 --- /dev/null +++ b/sdk/python/feast/templates/rag/bootstrap.py @@ -0,0 +1,56 @@ +def bootstrap(): + # Bootstrap() is called from init_repo() during `feast init` + import pathlib + from datetime import datetime + + import numpy as np + import pandas as pd + + repo_path = pathlib.Path(__file__).parent.absolute() / "feature_repo" + data_path = repo_path / "data" + data_path.mkdir(exist_ok=True) + + # Minimal city data with embeddings (384-d to match feature_store embedding_dim) + embedding_dim = 384 + now = datetime.now().replace(microsecond=0, tzinfo=None) + cities = [ + ( + 1, + "New York", + "New York", + "New York City is the most populous city in the United States.", + ), + ( + 2, + "Los Angeles", + "California", + "Los Angeles is the second most populous city in the United States.", + ), + ( + 3, + "Chicago", + "Illinois", + "Chicago is the third most populous city in the United States.", + ), + ] + rows = [] + for city_id, city_name, state, wiki_summary in cities: + vec = np.random.randn(embedding_dim).astype(np.float32) + vec = (vec / np.linalg.norm(vec)).tolist() + rows.append( + { + "city_id": city_id, + "event_timestamp": pd.Timestamp(now), + "vector": vec, + "sentence_chunks": wiki_summary[:200], + "state": f"{city_name}, {state}", + "wiki_summary": wiki_summary, + } + ) + df = pd.DataFrame(rows) + parquet_path = data_path / "city_wikipedia_summaries_with_embeddings.parquet" + df.to_parquet(path=str(parquet_path), index=False) + + +if __name__ == "__main__": + bootstrap() diff --git a/sdk/python/feast/templates/rag/feature_repo/__init__.py b/sdk/python/feast/templates/rag/feature_repo/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/sdk/python/feast/templates/rag/feature_repo/__init__.py @@ -0,0 +1 @@ + diff --git a/sdk/python/feast/templates/rag/feature_repo/example_repo.py b/sdk/python/feast/templates/rag/feature_repo/example_repo.py new file mode 100644 index 00000000000..af2ea96b1db --- /dev/null +++ b/sdk/python/feast/templates/rag/feature_repo/example_repo.py @@ -0,0 +1,137 @@ +from datetime import timedelta + +from feast import ( + Entity, + FeatureService, + FeatureView, + Field, + FileSource, + PushSource, +) +from feast.data_format import ParquetFormat +from feast.types import Array, Float32, String +from feast.value_type import ValueType + +# Entity: Identifies each city document/chunk in the knowledge base +city = Entity( + name="city_id", + value_type=ValueType.INT64, + description="Unique identifier for each city Wikipedia summary (document chunk ID).", + join_keys=["city_id"], +) + +# Data Source: Parquet file containing city summaries with pre-computed embeddings +city_summaries_source = FileSource( + name="city_summaries_source", + file_format=ParquetFormat(), + path="./data/city_wikipedia_summaries_with_embeddings.parquet", + timestamp_field="event_timestamp", + description="Wikipedia summaries of US cities (batch).", +) + +# Push Source: same schema as batch; allows near real-time ingestion of new/updated docs +city_summaries_push_source = PushSource( + name="city_summaries_push_source", + batch_source=city_summaries_source, + description="Push source for real-time updates to city summaries/embeddings.", +) + +# Feature View 1: City embeddings for semantic/vector search (RAG retrieval) +city_summary_embeddings = FeatureView( + name="city_summary_embeddings", + description="City Wikipedia summaries with embeddings for semantic search. ", + entities=[city], + schema=[ + Field( + name="vector", + dtype=Array(Float32), + description="384-dimensional sentence embedding for semantic similarity search (MiniLM).", + vector_index=True, + vector_search_metric="COSINE", + ), + Field( + name="sentence_chunks", + dtype=String, + description="Chunked sentences from the Wikipedia summary.", + ), + ], + source=city_summaries_source, + ttl=timedelta(days=1), + online=True, + tags={"team": "ml-platform", "use_case": "city_qa", "type": "vector"}, +) + +# Feature View 2: City metadata for scalar lookups (no vector search) +city_metadata = FeatureView( + name="city_metadata", + description="City metadata including state and full Wikipedia summary. ", + entities=[city], + schema=[ + Field( + name="state", + dtype=String, + description="US state where the city is located (e.g., 'New York, New York').", + ), + Field( + name="wiki_summary", + dtype=String, + description="Full Wikipedia summary of the city.", + ), + ], + source=city_summaries_source, + ttl=timedelta(hours=2), + online=True, + tags={"team": "ml-platform", "use_case": "city_qa", "type": "metadata"}, +) + +# Feature View 3: Fresh embeddings (PushSource) for near real-time doc updates +city_summary_embeddings_realtime = FeatureView( + name="city_summary_embeddings_realtime", + description="Same as city_summary_embeddings but with real-time ingestion (PushSource).", + entities=[city], + schema=[ + Field( + name="vector", + dtype=Array(Float32), + description="384-dimensional sentence embedding for semantic similarity search.", + vector_index=True, + vector_search_metric="COSINE", + ), + Field( + name="sentence_chunks", + dtype=String, + description="Chunked sentences from the Wikipedia summary.", + ), + ], + source=city_summaries_push_source, + ttl=timedelta(hours=2), + online=True, + tags={ + "team": "ml-platform", + "use_case": "city_qa", + "type": "vector", + "ingestion": "push", + }, +) + +# Feature Service: Bundles features for the City Q&A retrieval endpoint +city_qa_v1 = FeatureService( + name="city_qa_v1", + features=[ + city_summary_embeddings, + city_metadata, + ], + description="Feature service for City Information Q&A. ", + tags={"team": "ml-platform", "version": "v1"}, +) + +# Feature service that includes push-backed and request-time features +city_qa_v2 = FeatureService( + name="city_qa_v2", + features=[ + city_summary_embeddings_realtime, + city_metadata, + ], + description="City Q&A with push ingestion and request-time context (query_text, user_id).", + tags={"team": "ml-platform", "version": "v2"}, +) diff --git a/sdk/python/feast/templates/rag/feature_repo/feature_store.yaml b/sdk/python/feast/templates/rag/feature_repo/feature_store.yaml new file mode 100644 index 00000000000..658d710b247 --- /dev/null +++ b/sdk/python/feast/templates/rag/feature_repo/feature_store.yaml @@ -0,0 +1,20 @@ +project: my_project + +project_description: | + This project is a smart City Q&A assistant that gives you accurate, detailed answers about any city instantly. +provider: local +registry: data/registry.db +online_store: + type: milvus + path: data/online_store.db + vector_enabled: true + embedding_dim: 384 + index_type: "FLAT" + metric_type: "COSINE" + +offline_store: + type: file +entity_key_serialization_version: 3 +# By default, no_auth for authentication and authorization +auth: + type: no_auth diff --git a/sdk/python/feast/templates/rag/feature_repo/test_workflow.py b/sdk/python/feast/templates/rag/feature_repo/test_workflow.py new file mode 100644 index 00000000000..7a321208795 --- /dev/null +++ b/sdk/python/feast/templates/rag/feature_repo/test_workflow.py @@ -0,0 +1,155 @@ +from datetime import datetime + +import pandas as pd +import torch +import torch.nn.functional as F +from example_repo import ( + city, + city_metadata, + city_qa_v1, + city_qa_v2, + city_summary_embeddings, + city_summary_embeddings_realtime, +) +from transformers import AutoModel, AutoTokenizer + +from feast import FeatureStore +from feast.data_source import PushMode + +TOKENIZER = "sentence-transformers/all-MiniLM-L6-v2" +MODEL = "sentence-transformers/all-MiniLM-L6-v2" + + +def mean_pooling(model_output, attention_mask): + token_embeddings = model_output[0] + input_mask_expanded = ( + attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() + ) + return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp( + input_mask_expanded.sum(1), min=1e-9 + ) + + +def run_model(sentences, tokenizer, model): + encoded_input = tokenizer( + sentences, padding=True, truncation=True, return_tensors="pt" + ) + with torch.no_grad(): + model_output = model(**encoded_input) + sentence_embeddings = mean_pooling(model_output, encoded_input["attention_mask"]) + sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1) + return sentence_embeddings + + +def run_demo(): + print("City Information Q&A - RAG Demo") + + store = FeatureStore(repo_path=".") + print("\n[1/7] Applying feature definitions...") + store.apply( + [ + city, + city_summary_embeddings, + city_metadata, + city_summary_embeddings_realtime, + city_qa_v1, + city_qa_v2, + ] + ) + print( + " Entity, 3 feature views (batch + push), 2 feature services registered." + ) + + print("\n[2/7] Materializing batch views into online store...") + store.materialize( + feature_views=["city_summary_embeddings", "city_metadata"], + start_date=datetime(1970, 1, 1), + end_date=datetime.now(), + disable_event_timestamp=True, + ) + print(" city_summary_embeddings, city_metadata materialized from parquet.") + + print("\n[3/7] Verifying batch data...") + df = pd.read_parquet("./data/city_wikipedia_summaries_with_embeddings.parquet") + embedding_length = len(df["vector"][0]) + print(f" Parquet: {len(df)} rows, {embedding_length}-dim vectors.") + + tokenizer = AutoTokenizer.from_pretrained(TOKENIZER) + model = AutoModel.from_pretrained(MODEL) + + print("\n[4/7] Pushing one document to PushSource...") + push_text = "Demo pushed city for testing RAG and real-time ingestion." + push_embedding = run_model([push_text], tokenizer, model) + push_vector = push_embedding.detach().cpu().numpy().tolist()[0] + push_df = pd.DataFrame.from_dict( + { + "city_id": [99999], + "event_timestamp": [datetime.now()], + "vector": [push_vector], + "sentence_chunks": [push_text], + } + ) + store.push("city_summaries_push_source", push_df, to=PushMode.ONLINE) + print(" Pushed city_id=99999 to city_summary_embeddings_realtime.") + + print("\n[5/7] Vector search (batch view)...") + question = "the most populous city in the state of New York" + print(f' Query: "{question}"') + + query_embedding = run_model(question, tokenizer, model) + query = query_embedding.detach().cpu().numpy().tolist()[0] + + features_batch = store.retrieve_online_documents_v2( + features=[ + "city_summary_embeddings:vector", + "city_summary_embeddings:city_id", + "city_summary_embeddings:sentence_chunks", + ], + query=query, + top_k=3, + ) + print(" Top 3 (city_summary_embeddings):") + results_batch_df = features_batch.to_df() + print(results_batch_df[["city_id", "sentence_chunks", "distance"]].to_string()) + + print("\n[6/7] Vector search (realtime view, includes pushed doc)...") + question_realtime = "Demo pushed city for testing RAG" + query_realtime = ( + run_model([question_realtime], tokenizer, model) + .detach() + .cpu() + .numpy() + .tolist()[0] + ) + features_realtime = store.retrieve_online_documents_v2( + features=[ + "city_summary_embeddings_realtime:vector", + "city_summary_embeddings_realtime:city_id", + "city_summary_embeddings_realtime:sentence_chunks", + ], + query=query_realtime, + top_k=3, + ) + results_realtime_df = features_realtime.to_df() + print(" Top 3 (city_summary_embeddings_realtime):") + print(results_realtime_df[["city_id", "sentence_chunks", "distance"]].to_string()) + if any(results_realtime_df["city_id"] == 99999): + print(" Pushed doc (city_id=99999) in results.") + + print("\n[7/7] Metadata via Feature Service V2 (city_qa_v2)...") + top_city_id = int(results_batch_df["city_id"].iloc[0]) + + metadata_features = store.get_online_features( + features=store.get_feature_service("city_qa_v2"), + entity_rows=[{"city_id": top_city_id}], + ).to_dict() + + print(f" city_id={top_city_id} -> state: {metadata_features['state'][0]}") + print(f" wiki_summary: {metadata_features['wiki_summary'][0][:200]}...") + + store.teardown() + print("\nDone.") + + +if __name__ == "__main__": + run_demo() diff --git a/sdk/python/feast/templates/rag/gitignore b/sdk/python/feast/templates/rag/gitignore new file mode 100644 index 00000000000..2d7ceaab468 --- /dev/null +++ b/sdk/python/feast/templates/rag/gitignore @@ -0,0 +1,33 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*.pyo +*.pyd + +# C extensions +*.so + +# Distribution / packaging +.Python +env/ +venv/ +.venv + +# Pytest +.cache +.pytest_cache/ +.coverage +*.log + +# Jupyter Notebook +.ipynb_checkpoints + +# IDEs and Editors +.vscode/ +.idea/ + +# OS generated files +.DS_Store + +# Feast +.feast/