feat: Adding optional embeddings to content

Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>

tests/integration/vector_io/test_openai_vector_stores.py

@@ -11,6 +11,7 @@ import pytest
 from llama_stack_client import BadRequestError
 from openai import BadRequestError as OpenAIBadRequestError
 
+from llama_stack.apis.files import ExpiresAfter
 from llama_stack.apis.vector_io import Chunk
 from llama_stack.core.library_client import LlamaStackAsLibraryClient
 from llama_stack.log import get_logger
@@ -1604,3 +1605,97 @@ def test_openai_vector_store_embedding_config_from_metadata(
 
     assert "metadata_config_store" in store_names
     assert "consistent_config_store" in store_names
+
+
+@vector_provider_wrapper
+def test_openai_vector_store_file_contents_with_extra_query(
+    compat_client_with_empty_stores, client_with_models, embedding_model_id, embedding_dimension, vector_io_provider_id
+):
+    """Test that vector store file contents endpoint supports extra_query parameter."""
+    skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
+    compat_client = compat_client_with_empty_stores
+
+    # Create a vector store
+    vector_store = compat_client.vector_stores.create(
+        name="test_extra_query_store",
+        extra_body={
+            "embedding_model": embedding_model_id,
+            "provider_id": vector_io_provider_id,
+        },
+    )
+
+    # Create and attach a file
+    test_content = b"This is test content for extra_query validation."
+    with BytesIO(test_content) as file_buffer:
+        file_buffer.name = "test_extra_query.txt"
+        file = compat_client.files.create(
+            file=file_buffer,
+            purpose="assistants",
+            expires_after=ExpiresAfter(anchor="created_at", seconds=86400),
+        )
+
+    file_attach_response = compat_client.vector_stores.files.create(
+        vector_store_id=vector_store.id,
+        file_id=file.id,
+        extra_body={"embedding_model": embedding_model_id},
+    )
+    assert file_attach_response.status == "completed"
+
+    # Wait for processing
+    time.sleep(2)
+
+    # Test that extra_query parameter is accepted and processed
+    content_with_extra_query = compat_client.vector_stores.files.content(
+        vector_store_id=vector_store.id,
+        file_id=file.id,
+        extra_query={"include_embeddings": True, "include_metadata": True},
+    )
+
+    # Test without extra_query for comparison
+    content_without_extra_query = compat_client.vector_stores.files.content(
+        vector_store_id=vector_store.id,
+        file_id=file.id,
+    )
+
+    # Validate that both calls succeed
+    assert content_with_extra_query is not None
+    assert content_without_extra_query is not None
+    assert len(content_with_extra_query.content) > 0
+    assert len(content_without_extra_query.content) > 0
+
+    # Validate that extra_query parameter is processed correctly
+    # Both should have the embedding/metadata fields available (may be None based on flags)
+    first_chunk_with_flags = content_with_extra_query.content[0]
+    first_chunk_without_flags = content_without_extra_query.content[0]
+
+    # The key validation: extra_query fields are present in the response
+    # Handle both dict and object responses (different clients may return different formats)
+    def has_field(obj, field):
+        if isinstance(obj, dict):
+            return field in obj
+        else:
+            return hasattr(obj, field)
+
+    # Validate that all expected fields are present in both responses
+    expected_fields = ["embedding", "chunk_metadata", "metadata", "text"]
+    for field in expected_fields:
+        assert has_field(first_chunk_with_flags, field), f"Field '{field}' missing from response with extra_query"
+        assert has_field(first_chunk_without_flags, field), f"Field '{field}' missing from response without extra_query"
+
+    # Validate content is the same
+    def get_field(obj, field):
+        if isinstance(obj, dict):
+            return obj[field]
+        else:
+            return getattr(obj, field)
+
+    assert get_field(first_chunk_with_flags, "text") == test_content.decode("utf-8")
+    assert get_field(first_chunk_without_flags, "text") == test_content.decode("utf-8")
+
+    with_flags_embedding = get_field(first_chunk_with_flags, "embedding")
+    without_flags_embedding = get_field(first_chunk_without_flags, "embedding")
+
+    # Validate that embeddings are included when requested and excluded when not requested
+    assert with_flags_embedding is not None, "Embeddings should be included when include_embeddings=True"
+    assert len(with_flags_embedding) > 0, "Embedding should be a non-empty list"
+    assert without_flags_embedding is None, "Embeddings should not be included when include_embeddings=False"

tests/unit/core/routers/test_vector_io.py

@@ -55,3 +55,65 @@ async def test_create_vector_stores_multiple_providers_missing_provider_id_error
 
     with pytest.raises(ValueError, match="Multiple vector_io providers available"):
         await router.openai_create_vector_store(request)
+
+
+async def test_update_vector_store_provider_id_change_fails():
+    """Test that updating a vector store with a different provider_id fails with clear error."""
+    mock_routing_table = Mock()
+
+    # Mock an existing vector store with provider_id "faiss"
+    mock_existing_store = Mock()
+    mock_existing_store.provider_id = "inline::faiss"
+    mock_existing_store.identifier = "vs_123"
+
+    mock_routing_table.get_object_by_identifier = AsyncMock(return_value=mock_existing_store)
+    mock_routing_table.get_provider_impl = AsyncMock(
+        return_value=Mock(openai_update_vector_store=AsyncMock(return_value=Mock(id="vs_123")))
+    )
+
+    router = VectorIORouter(mock_routing_table)
+
+    # Try to update with different provider_id in metadata - this should fail
+    with pytest.raises(ValueError, match="provider_id cannot be changed after vector store creation"):
+        await router.openai_update_vector_store(
+            vector_store_id="vs_123",
+            name="updated_name",
+            metadata={"provider_id": "inline::sqlite"},  # Different provider_id
+        )
+
+    # Verify the existing store was looked up to check provider_id
+    mock_routing_table.get_object_by_identifier.assert_called_once_with("vector_store", "vs_123")
+
+    # Provider should not be called since validation failed
+    mock_routing_table.get_provider_impl.assert_not_called()
+
+
+async def test_update_vector_store_same_provider_id_succeeds():
+    """Test that updating a vector store with the same provider_id succeeds."""
+    mock_routing_table = Mock()
+
+    # Mock an existing vector store with provider_id "faiss"
+    mock_existing_store = Mock()
+    mock_existing_store.provider_id = "inline::faiss"
+    mock_existing_store.identifier = "vs_123"
+
+    mock_routing_table.get_object_by_identifier = AsyncMock(return_value=mock_existing_store)
+    mock_routing_table.get_provider_impl = AsyncMock(
+        return_value=Mock(openai_update_vector_store=AsyncMock(return_value=Mock(id="vs_123")))
+    )
+
+    router = VectorIORouter(mock_routing_table)
+
+    # Update with same provider_id should succeed
+    await router.openai_update_vector_store(
+        vector_store_id="vs_123",
+        name="updated_name",
+        metadata={"provider_id": "inline::faiss"},  # Same provider_id
+    )
+
+    # Verify the provider update method was called
+    mock_routing_table.get_provider_impl.assert_called_once_with("vs_123")
+    provider = await mock_routing_table.get_provider_impl("vs_123")
+    provider.openai_update_vector_store.assert_called_once_with(
+        vector_store_id="vs_123", name="updated_name", expires_after=None, metadata={"provider_id": "inline::faiss"}
+    )

tests/unit/server/test_query_params_middleware.py

@@ -0,0 +1,86 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from unittest.mock import AsyncMock, Mock
+
+from fastapi import Request
+
+from llama_stack.core.server.query_params_middleware import QueryParamsMiddleware
+
+
+class TestQueryParamsMiddleware:
+    """Test cases for the QueryParamsMiddleware."""
+
+    async def test_extracts_query_params_for_vector_store_content(self):
+        """Test that middleware extracts query params for vector store content endpoints."""
+        middleware = QueryParamsMiddleware(Mock())
+        request = Mock(spec=Request)
+        request.method = "GET"
+
+        # Mock the URL properly
+        mock_url = Mock()
+        mock_url.path = "/v1/vector_stores/vs_123/files/file_456/content"
+        request.url = mock_url
+
+        request.query_params = {"include_embeddings": "true", "include_metadata": "false"}
+
+        # Create a fresh state object without any attributes
+        class MockState:
+            pass
+
+        request.state = MockState()
+
+        await middleware.dispatch(request, AsyncMock())
+
+        assert hasattr(request.state, "extra_query")
+        assert request.state.extra_query == {"include_embeddings": True, "include_metadata": False}
+
+    async def test_ignores_non_vector_store_endpoints(self):
+        """Test that middleware ignores non-vector store endpoints."""
+        middleware = QueryParamsMiddleware(Mock())
+        request = Mock(spec=Request)
+        request.method = "GET"
+
+        # Mock the URL properly
+        mock_url = Mock()
+        mock_url.path = "/v1/inference/chat_completion"
+        request.url = mock_url
+
+        request.query_params = {"include_embeddings": "true"}
+
+        # Create a fresh state object without any attributes
+        class MockState:
+            pass
+
+        request.state = MockState()
+
+        await middleware.dispatch(request, AsyncMock())
+
+        assert not hasattr(request.state, "extra_query")
+
+    async def test_handles_json_parsing(self):
+        """Test that middleware correctly parses JSON values and handles invalid JSON."""
+        middleware = QueryParamsMiddleware(Mock())
+        request = Mock(spec=Request)
+        request.method = "GET"
+
+        # Mock the URL properly
+        mock_url = Mock()
+        mock_url.path = "/v1/vector_stores/vs_123/files/file_456/content"
+        request.url = mock_url
+
+        request.query_params = {"config": '{"key": "value"}', "invalid": "not-json{", "number": "42"}
+
+        # Create a fresh state object without any attributes
+        class MockState:
+            pass
+
+        request.state = MockState()
+
+        await middleware.dispatch(request, AsyncMock())
+
+        expected = {"config": {"key": "value"}, "invalid": "not-json{", "number": 42}
+        assert request.state.extra_query == expected

tests/unit/server/test_sse.py

@@ -104,12 +104,18 @@ async def test_paginated_response_url_setting():
 
     route_handler = create_dynamic_typed_route(mock_api_method, "get", "/test/route")
 
-    # Mock minimal request
+    # Mock minimal request with proper state object
     request = MagicMock()
     request.scope = {"user_attributes": {}, "principal": ""}
     request.headers = {}
     request.body = AsyncMock(return_value=b"")
 
+    # Create a simple state object without auto-generating attributes
+    class MockState:
+        pass
+
+    request.state = MockState()
+
     result = await route_handler(request)
 
     assert isinstance(result, PaginatedResponse)