Added test for testing the new OpenAI api

Signed-off-by: Chaitany patel <[email protected]>

Author: patelchaitany

sdk/python/tests/integration/online_store/test_universal_online.py

@@ -14,13 +14,15 @@
 
 from feast import FeatureStore
 from feast.entity import Entity
-from feast.errors import FeatureNameCollisionError
+from feast.errors import FeatureNameCollisionError, FeatureViewNotFoundException
 from feast.feature_service import FeatureService
 from feast.feature_view import FeatureView
 from feast.field import Field
+from feast.filter_models import ComparisonFilter, CompoundFilter
 from feast.infra.offline_stores.file_source import FileSource
 from feast.infra.utils.postgres.postgres_config import ConnectionType
 from feast.online_response import TIMESTAMP_POSTFIX
+from feast.repo_config import EmbeddingModelConfig
 from feast.types import (
     Array,
     Float32,
@@ -1265,3 +1267,313 @@ def test_retrieve_online_documents_v2(environment, fake_document_data):
     assert len(no_match_results["text_field"]) == 0
     assert "text_rank" in no_match_results
     assert len(no_match_results["text_rank"]) == 0
+
+
+def _setup_documents_with_categories(fs):
+    """Shared helper that creates and populates a feature view with embeddings,
+    text, and category fields. Returns (feature_view, entity, dataframe)."""
+    n_rows = 20
+    vector_dim = 2
+    random.seed(42)
+
+    df = pd.DataFrame(
+        {
+            "item_id": list(range(n_rows)),
+            "embedding": [list(np.random.random(vector_dim)) for _ in range(n_rows)],
+            "text_field": [
+                f"Document text content {i} with searchable keywords"
+                for i in range(n_rows)
+            ],
+            "category": [f"Category-{i % 5}" for i in range(n_rows)],
+            "event_timestamp": [datetime.now() for _ in range(n_rows)],
+        }
+    )
+
+    data_source = FileSource(
+        path="dummy_path.parquet", timestamp_field="event_timestamp"
+    )
+
+    item_entity = Entity(
+        name="item_id",
+        join_keys=["item_id"],
+        value_type=ValueType.INT64,
+    )
+
+    item_embeddings_fv = FeatureView(
+        name="item_embeddings",
+        entities=[item_entity],
+        schema=[
+            Field(name="embedding", dtype=Array(Float32), vector_index=True),
+            Field(name="text_field", dtype=String),
+            Field(name="category", dtype=String),
+            Field(name="item_id", dtype=Int64),
+        ],
+        source=data_source,
+    )
+
+    fs.apply([item_embeddings_fv, item_entity])
+    fs.write_to_online_store("item_embeddings", df)
+    return item_embeddings_fv, item_entity, df
+
+
+@pytest.mark.integration
+@pytest.mark.universal_online_stores(only=["pgvector", "elasticsearch"])
+def test_retrieve_online_documents_v2_with_filters(environment, fake_document_data):
+    """Test that metadata filters narrow down vector/text search results."""
+    fs = environment.feature_store
+    fs.config.online_store.vector_enabled = True
+
+    _, _, df = _setup_documents_with_categories(fs)
+    vector_dim = 2
+    query_embedding = list(np.random.random(vector_dim))
+
+    # --- eq filter: only Category-0 rows ---
+    eq_filter = ComparisonFilter(type="eq", key="category", value="Category-0")
+    results = fs.retrieve_online_documents_v2(
+        features=[
+            "item_embeddings:embedding",
+            "item_embeddings:text_field",
+            "item_embeddings:category",
+            "item_embeddings:item_id",
+        ],
+        query=query_embedding,
+        top_k=10,
+        distance_metric="L2",
+        filters=eq_filter,
+    ).to_dict()
+
+    assert len(results["category"]) > 0
+    assert len(results["category"]) <= 4  # 20 rows / 5 categories
+    assert all(c == "Category-0" for c in results["category"])
+
+    # --- ne filter: exclude Category-0 ---
+    ne_filter = ComparisonFilter(type="ne", key="category", value="Category-0")
+    results = fs.retrieve_online_documents_v2(
+        features=[
+            "item_embeddings:embedding",
+            "item_embeddings:text_field",
+            "item_embeddings:category",
+            "item_embeddings:item_id",
+        ],
+        query=query_embedding,
+        top_k=10,
+        distance_metric="L2",
+        filters=ne_filter,
+    ).to_dict()
+
+    assert len(results["category"]) > 0
+    assert all(c != "Category-0" for c in results["category"])
+
+    # --- in filter: Category-0 or Category-1 ---
+    in_filter = ComparisonFilter(
+        type="in", key="category", value=["Category-0", "Category-1"]
+    )
+    results = fs.retrieve_online_documents_v2(
+        features=[
+            "item_embeddings:embedding",
+            "item_embeddings:text_field",
+            "item_embeddings:category",
+            "item_embeddings:item_id",
+        ],
+        query=query_embedding,
+        top_k=10,
+        distance_metric="L2",
+        filters=in_filter,
+    ).to_dict()
+
+    assert len(results["category"]) > 0
+    assert all(c in ("Category-0", "Category-1") for c in results["category"])
+
+    # --- compound AND filter: category == Category-0 AND item_id >= 5 ---
+    and_filter = CompoundFilter(
+        type="and",
+        filters=[
+            ComparisonFilter(type="eq", key="category", value="Category-0"),
+            ComparisonFilter(type="gte", key="item_id", value=5),
+        ],
+    )
+    results = fs.retrieve_online_documents_v2(
+        features=[
+            "item_embeddings:embedding",
+            "item_embeddings:text_field",
+            "item_embeddings:category",
+            "item_embeddings:item_id",
+        ],
+        query=query_embedding,
+        top_k=10,
+        distance_metric="L2",
+        filters=and_filter,
+    ).to_dict()
+
+    assert len(results["category"]) > 0
+    assert all(c == "Category-0" for c in results["category"])
+    assert all(i >= 5 for i in results["item_id"])
+
+    # --- text search + filter ---
+    text_filter = ComparisonFilter(type="eq", key="category", value="Category-2")
+    text_results = fs.retrieve_online_documents_v2(
+        features=[
+            "item_embeddings:embedding",
+            "item_embeddings:text_field",
+            "item_embeddings:category",
+            "item_embeddings:item_id",
+        ],
+        query_string="searchable keywords",
+        top_k=10,
+        filters=text_filter,
+    ).to_dict()
+
+    assert len(text_results["category"]) > 0
+    assert all(c == "Category-2" for c in text_results["category"])
+
+    # --- filter with no matches ---
+    empty_filter = ComparisonFilter(
+        type="eq", key="category", value="NonexistentCategory"
+    )
+    empty_results = fs.retrieve_online_documents_v2(
+        features=[
+            "item_embeddings:embedding",
+            "item_embeddings:text_field",
+            "item_embeddings:category",
+            "item_embeddings:item_id",
+        ],
+        query=query_embedding,
+        top_k=10,
+        distance_metric="L2",
+        filters=empty_filter,
+    ).to_dict()
+
+    assert len(empty_results.get("category", [])) == 0
+
+
+@pytest.mark.integration
+@pytest.mark.universal_online_stores(only=["pgvector", "elasticsearch"])
+def test_retrieve_online_documents_openai(environment, fake_document_data):
+    """Test OpenAI-compatible vector store search returns the correct response shape."""
+    fs = environment.feature_store
+    fs.config.online_store.vector_enabled = True
+
+    fv, _, df = _setup_documents_with_categories(fs)
+    vector_dim = 2
+
+    fs.config.embedding_model = EmbeddingModelConfig(model="text-embedding-3-small")
+
+    fake_embedding = list(np.random.random(vector_dim))
+    mock_embed_response = unittest.mock.MagicMock()
+    mock_embed_response.data = [{"embedding": fake_embedding}]
+
+    with unittest.mock.patch(
+        "feast.feature_store.litellm_embedding", create=True
+    ) as mock_litellm:
+        mock_litellm.return_value = mock_embed_response
+
+        with unittest.mock.patch(
+            "feast.feature_store.FeatureStore.retrieve_online_documents_openai",
+            wraps=fs.retrieve_online_documents_openai,
+        ):
+            # Patch the litellm import inside the method
+            with unittest.mock.patch.dict(
+                "sys.modules",
+                {"litellm": unittest.mock.MagicMock(embedding=mock_litellm)},
+            ):
+                result = fs.retrieve_online_documents_openai(
+                    vector_store_id="item_embeddings",
+                    query="test query",
+                    max_num_results=5,
+                )
+
+    # Validate top-level OpenAI response shape
+    assert result["object"] == "vector_store.search_results.page"
+    assert isinstance(result["search_query"], list)
+    assert result["search_query"] == ["test query"]
+    assert result["has_more"] is False
+    assert result["next_page"] is None
+
+    assert isinstance(result["data"], list)
+    assert len(result["data"]) > 0
+    assert len(result["data"]) <= 5
+
+    for item_result in result["data"]:
+        assert "file_id" in item_result
+        assert "filename" in item_result
+        assert item_result["filename"] == "item_embeddings"
+        assert "score" in item_result
+        assert isinstance(item_result["score"], float)
+        assert "attributes" in item_result
+        assert isinstance(item_result["attributes"], dict)
+        assert "content" in item_result
+        assert isinstance(item_result["content"], list)
+        for part in item_result["content"]:
+            assert "type" in part
+            assert part["type"] == "text"
+            assert "text" in part
+
+    # --- Test with features_to_retrieve ---
+    with unittest.mock.patch.dict(
+        "sys.modules",
+        {
+            "litellm": unittest.mock.MagicMock(
+                embedding=unittest.mock.MagicMock(return_value=mock_embed_response)
+            ),
+        },
+    ):
+        result_subset = fs.retrieve_online_documents_openai(
+            vector_store_id="item_embeddings",
+            query="test query",
+            max_num_results=5,
+            features_to_retrieve=["text_field", "category"],
+        )
+
+    assert len(result_subset["data"]) > 0
+    for item_result in result_subset["data"]:
+        attr_keys = set(item_result["attributes"].keys())
+        assert "embedding" not in attr_keys
+
+    # --- Test with list query ---
+    with unittest.mock.patch.dict(
+        "sys.modules",
+        {
+            "litellm": unittest.mock.MagicMock(
+                embedding=unittest.mock.MagicMock(return_value=mock_embed_response)
+            ),
+        },
+    ):
+        result_list = fs.retrieve_online_documents_openai(
+            vector_store_id="item_embeddings",
+            query=["term1", "term2"],
+            max_num_results=5,
+        )
+
+    assert result_list["search_query"] == ["term1", "term2"]
+
+
+@pytest.mark.integration
+@pytest.mark.universal_online_stores(only=["pgvector", "elasticsearch"])
+def test_retrieve_online_documents_openai_no_embedding_config(
+    environment, fake_document_data
+):
+    """Test that retrieve_online_documents_openai raises ValueError
+    when embedding_model is not configured."""
+    fs = environment.feature_store
+    fs.config.embedding_model = None
+
+    with pytest.raises(ValueError, match="embedding_model is not configured"):
+        fs.retrieve_online_documents_openai(
+            vector_store_id="item_embeddings",
+            query="test query",
+        )
+
+
+@pytest.mark.integration
+@pytest.mark.universal_online_stores(only=["pgvector", "elasticsearch"])
+def test_retrieve_online_documents_openai_not_found(environment, fake_document_data):
+    """Test that retrieve_online_documents_openai raises FeatureViewNotFoundException
+    for a non-existent feature view."""
+    fs = environment.feature_store
+    fs.config.embedding_model = EmbeddingModelConfig(model="text-embedding-3-small")
+
+    with pytest.raises(FeatureViewNotFoundException):
+        fs.retrieve_online_documents_openai(
+            vector_store_id="nonexistent_feature_view",
+            query="test query",
+        )