feat: Add feast rag retriver functionality (feast-dev#5405)

Add feast rag retriver functionality

Signed-off-by: Fiona Waters <[email protected]>

Author: Fiona-Waters

.github/workflows/unit_tests.yml

@@ -30,12 +30,6 @@ jobs:
         uses: astral-sh/setup-uv@v5
         with:
           enable-cache: true
-      - name: Install torch (platform-specific)
-        run: |
-          if [[ "$RUNNER_OS" == "Linux" ]]; then
-            pip install torch torchvision \
-              --index-url https://download.pytorch.org/whl/cpu
-          fi
       - name: Install dependencies
         run: make install-python-dependencies-ci
       - name: Test Python

Makefile

@@ -81,6 +81,9 @@ install-python-dependencies-minimal: ## Install minimal Python dependencies usin
 # Used in github actions/ci
 # formerly install-python-ci-dependencies-uv
 install-python-dependencies-ci: ## Install Python CI dependencies in system environment using uv
+	# Install CPU-only torch first to prevent CUDA dependency issues
+	pip uninstall torch torchvision -y || true
+	pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu --force-reinstall
 	uv pip sync --system sdk/python/requirements/py$(PYTHON_VERSION)-ci-requirements.txt
 	uv pip install --system --no-deps -e .
 

examples/rag-retriever/README.md

@@ -0,0 +1,72 @@
+# End-to-end RAG example using Feast and Milvus.
+
+## Introduction
+This example notebook provides a step-by-step demonstration of building and using a RAG system with Feast Feature Store and the custom FeastRagRetriever. The notebook walks through:
+
+1. Data Preparation
+   - Loads a subset of the Wikipedia DPR dataset (1% of training data)
+   - Implements text chunking with configurable chunk size and overlap
+   - Processes text into manageable passages with unique IDs
+
+2. Embedding Generation
+   - Uses `all-MiniLM-L6-v2` sentence transformer model
+   - Generates 384-dimensional embeddings for text passages
+   - Demonstrates batch processing with GPU support
+
+3. Feature Store Setup
+   - Creates a Parquet file as the historical data source
+   - Configures Feast with the feature repository
+   - Demonstrates writing embeddings from data source to Milvus online store which can be used for model training later
+
+4. RAG System Implementation
+   - **Embedding Model**: `all-MiniLM-L6-v2` (configurable)
+   - **Generator Model**: `granite-3.2-2b-instruct` (configurable)
+   - **Vector Store**: Custom implementation with Feast integration
+   - **Retriever**: Custom implementation extending HuggingFace's RagRetriever
+
+5. Query Demonstration
+   - Perform inference with retrieved context
+
+## Requirements
+ - A Kubernetes cluster with:
+   - GPU nodes available (for model inference)
+   - At least 200GB of storage
+   - A standalone Milvus deployment. See example [here](https://github.com/milvus-io/milvus-helm/tree/master/charts/milvus).
+
+## Running the example
+Clone this repository: https://github.com/feast-dev/feast.git
+Navigate to the examples/rag-retriever directory. Here you will find the following files: 
+
+* **feature_repo/feature_store.yaml**
+  This is the core configuration file for the RAG project's feature store, configuring a Milvus online store on a local provider. 
+  * In order to configure Milvus you should:
+     - Update `feature_store.yaml` with your Milvus connection details:
+       - host
+       - port (default: 19530)
+       - credentials (if required)
+
+* **__feature_repo/ragproject_repo.py__**
+  This is the Feast feature repository configuration that defines the schema and data source for Wikipedia passage embeddings. 
+
+* **__rag_feast.ipynb__**
+  This is a notebook demonstrating the implementation of a RAG system using Feast feature store. The notebook provides:
+
+  - A complete end-to-end example of building a RAG system with:
+    - Data preparation using the Wiki DPR dataset
+    - Text chunking and preprocessing
+    - Vector embedding generation using sentence-transformers
+    - Integration with Milvus vector store
+    - Inference utilising a custom RagRetriever: FeastRagRetriever
+  - Uses `all-MiniLM-L6-v2` for generating embeddings
+  - Implements `granite-3.2-2b-instruct` as the generator model
+
+Open `rag_feast.ipynb` and follow the steps in the notebook to run the example.
+
+## FeastRagRetriver Low Level Design
+
+<img src="images/FeastRagRetriever.png" width="800" height="450" alt="Low level design for feast rag retriever">
+
+## Helpful Information
+- Ensure your Milvus instance is properly configured and running
+- Vector dimensions and similarity metrics can be adjusted in the feature store configuration
+- The example uses Wikipedia data, but the system can be adapted for other datasets

examples/rag-retriever/feature_repo/feature_store.yaml

@@ -0,0 +1,18 @@
+project: ragproject
+provider: local
+registry: data/registry.db
+online_store:
+  type: milvus
+  host: # Insert Milvus route host
+  username: # Insert Milvus username if required
+  password: # Insert Milvus password if required
+  port: 19530
+  vector_enabled: true
+  embedding_dim: 384
+  index_type: FLAT
+  metric_type: COSINE
+offline_store:
+  type: file
+entity_key_serialization_version: 3
+auth:
+    type: no_auth

examples/rag-retriever/feature_repo/ragproject_repo.py

@@ -0,0 +1,48 @@
+from datetime import timedelta
+
+from feast import Entity, FeatureView, Field, FileSource, ValueType
+from feast.data_format import ParquetFormat
+from feast.types import Array, Float32, String
+
+# Define your entity (primary key for feature lookup)
+wiki_passage = Entity(
+    name="passage_id",
+    join_keys=["passage_id"],
+    value_type=ValueType.STRING,
+    description="Unique ID of a Wikipedia passage",
+)
+
+parquet_file_path = "data/wiki_dpr.parquet"
+
+# Define offline source
+wiki_dpr_source = FileSource(
+    name="wiki_dpr_source",
+    file_format=ParquetFormat(),
+    path=parquet_file_path,
+    timestamp_field="event_timestamp",
+)
+
+# Define the feature view for the Wikipedia passage content
+wiki_passage_feature_view = FeatureView(
+    name="wiki_passages",
+    entities=[wiki_passage],
+    ttl=timedelta(days=1),
+    schema=[
+        Field(
+            name="passage_text",
+            dtype=String,
+            description="Content of the Wikipedia passage",
+        ),
+        Field(
+            name="embedding",
+            dtype=Array(Float32),
+            description="vectors",
+            vector_index=True,
+            vector_length=384,
+            vector_search_metric="COSINE",
+        ),
+    ],
+    online=True,
+    source=wiki_dpr_source,
+    description="Content features of Wikipedia passages",
+)