Add RAG project with feature views and permissions

Co-Authored-By: Francisco Javier Arceo <arceofrancisco@gmail.com>

ui/feature_repo/apply_permissions.py

@@ -4,6 +4,9 @@
     zipcode_source_permission,
     model_v1_permission,
     risky_features_permission,
+    document_embeddings_permission,
+    document_metadata_permission,
+    rag_model_permission,
 )
 
 store = FeatureStore(repo_path=".")
@@ -13,6 +16,9 @@
     zipcode_source_permission,
     model_v1_permission,
     risky_features_permission,
+    document_embeddings_permission,
+    document_metadata_permission,
+    rag_model_permission,
 ])
 
 print("Permissions applied successfully!")

ui/feature_repo/features.py

@@ -1,14 +1,15 @@
 from datetime import timedelta
 
 import pandas as pd
+import numpy as np
 
 from feast import Entity, FeatureService, FeatureView, Field, FileSource
 from feast.data_source import RequestSource
 from feast.on_demand_feature_view import on_demand_feature_view
 from feast.permissions.action import AuthzedAction, READ
 from feast.permissions.permission import Permission
 from feast.permissions.policy import RoleBasedPolicy
-from feast.types import Bool, Int64, String
+from feast.types import Bool, Int64, String, Float32, Array
 
 zipcode = Entity(
     name="zipcode",
@@ -234,3 +235,121 @@ def transaction_gt_last_credit_card_due(inputs: pd.DataFrame) -> pd.DataFrame:
     policy=RoleBasedPolicy(roles=["trusted_analyst"]),
     actions=[AuthzedAction.READ_OFFLINE],
 )
+
+document = Entity(
+    name="document_id",
+    description="Document identifier for RAG system",
+    tags={
+        "owner": "nlp_team@feast.ai",
+        "team": "rag",
+    },
+)
+
+document_source = FileSource(
+    name="document_embeddings",
+    path="data/document_embeddings.parquet",
+    timestamp_field="event_timestamp",
+    created_timestamp_column="created_timestamp",
+)
+
+document_metadata_source = FileSource(
+    name="document_metadata",
+    path="data/document_metadata.parquet",
+    timestamp_field="event_timestamp",
+    created_timestamp_column="created_timestamp",
+)
+
+document_embeddings_view = FeatureView(
+    name="document_embeddings",
+    entities=[document],
+    ttl=timedelta(days=365),
+    schema=[
+        Field(name="embedding", dtype=Array(Float32, 768)),
+        Field(name="document_id", dtype=String),
+    ],
+    source=document_source,
+    tags={
+        "date_added": "2025-05-04",
+        "model": "sentence-transformer",
+        "access_group": "nlp-team@feast.ai",
+        "stage": "prod",
+    },
+    online=True,
+)
+
+document_metadata_view = FeatureView(
+    name="document_metadata",
+    entities=[document],
+    ttl=timedelta(days=365),
+    schema=[
+        Field(name="title", dtype=String),
+        Field(name="content", dtype=String),
+        Field(name="source", dtype=String),
+        Field(name="author", dtype=String),
+        Field(name="publish_date", dtype=String),
+        Field(name="document_id", dtype=String),
+    ],
+    source=document_metadata_source,
+    tags={
+        "date_added": "2025-05-04",
+        "access_group": "nlp-team@feast.ai",
+        "stage": "prod",
+    },
+    online=True,
+)
+
+# Define a request data source for query embeddings
+query_request = RequestSource(
+    name="query",
+    schema=[
+        Field(name="query_embedding", dtype=Array(Float32, 768)),
+    ],
+)
+
+# Define an on-demand feature view for similarity calculation
+@on_demand_feature_view(
+    sources=[document_embeddings_view, query_request],
+    schema=[
+        Field(name="similarity_score", dtype=Float32),
+    ],
+)
+def document_similarity(inputs: pd.DataFrame) -> pd.DataFrame:
+    """Calculate cosine similarity between query and document embeddings."""
+    df = pd.DataFrame()
+    df["similarity_score"] = 0.95  # Placeholder value
+    return df
+
+rag_model = FeatureService(
+    name="rag_retriever",
+    features=[
+        document_embeddings_view,
+        document_metadata_view,
+        document_similarity,
+    ],
+    tags={"owner": "nlp_team@feast.ai", "stage": "prod"},
+    description="Retrieval Augmented Generation model",
+)
+
+document_embeddings_permission = Permission(
+    name="document-embeddings-reader",
+    types=[FeatureView],
+    name_patterns=["document_embeddings"],
+    policy=RoleBasedPolicy(roles=["ml_engineer", "data_scientist"]),
+    actions=[AuthzedAction.DESCRIBE, *READ],
+)
+
+document_metadata_permission = Permission(
+    name="document-metadata-reader",
+    types=[FeatureView],
+    name_patterns=["document_metadata"],
+    policy=RoleBasedPolicy(roles=["ml_engineer", "content_manager"]),
+    actions=[AuthzedAction.DESCRIBE, *READ],
+)
+
+rag_model_permission = Permission(
+    name="rag-model-user",
+    types=[FeatureService],
+    name_patterns=["rag_retriever"],
+    policy=RoleBasedPolicy(roles=["ml_engineer", "app_developer"]),
+    actions=[AuthzedAction.DESCRIBE, AuthzedAction.READ_ONLINE],
+)