feat: Updating files/content response to return additional fields

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

tests/unit/providers/vector_io/remote/test_milvus.py

@@ -15,6 +15,9 @@ from llama_stack.apis.vector_io import QueryChunksResponse
 pymilvus_mock = MagicMock()
 pymilvus_mock.DataType = MagicMock()
 pymilvus_mock.MilvusClient = MagicMock
+pymilvus_mock.RRFRanker = MagicMock
+pymilvus_mock.WeightedRanker = MagicMock
+pymilvus_mock.AnnSearchRequest = MagicMock
 
 # Apply the mock before importing MilvusIndex
 with patch.dict("sys.modules", {"pymilvus": pymilvus_mock}):
@@ -183,3 +186,141 @@ async def test_delete_collection(milvus_index, mock_milvus_client):
     await milvus_index.delete()
 
     mock_milvus_client.drop_collection.assert_called_once_with(collection_name=milvus_index.collection_name)
+
+
+async def test_query_hybrid_search_rrf(
+    milvus_index, sample_chunks, sample_embeddings, embedding_dimension, mock_milvus_client
+):
+    """Test hybrid search with RRF reranker."""
+    mock_milvus_client.has_collection.return_value = True
+    await milvus_index.add_chunks(sample_chunks, sample_embeddings)
+
+    # Mock hybrid search results
+    mock_milvus_client.hybrid_search.return_value = [
+        [
+            {
+                "id": 0,
+                "distance": 0.1,
+                "entity": {"chunk_content": {"content": "mock chunk 1", "metadata": {"document_id": "doc1"}}},
+            },
+            {
+                "id": 1,
+                "distance": 0.2,
+                "entity": {"chunk_content": {"content": "mock chunk 2", "metadata": {"document_id": "doc2"}}},
+            },
+        ]
+    ]
+
+    # Test hybrid search with RRF reranker
+    query_embedding = np.random.rand(embedding_dimension).astype(np.float32)
+    query_string = "test query"
+    response = await milvus_index.query_hybrid(
+        embedding=query_embedding,
+        query_string=query_string,
+        k=2,
+        score_threshold=0.0,
+        reranker_type="rrf",
+        reranker_params={"impact_factor": 60.0},
+    )
+
+    assert isinstance(response, QueryChunksResponse)
+    assert len(response.chunks) == 2
+    assert len(response.scores) == 2
+
+    # Verify hybrid search was called with correct parameters
+    mock_milvus_client.hybrid_search.assert_called_once()
+    call_args = mock_milvus_client.hybrid_search.call_args
+
+    # Check that the request contains both vector and BM25 search requests
+    reqs = call_args[1]["reqs"]
+    assert len(reqs) == 2
+    assert reqs[0].anns_field == "vector"
+    assert reqs[1].anns_field == "sparse"
+    ranker = call_args[1]["ranker"]
+    assert ranker is not None
+
+
+async def test_query_hybrid_search_weighted(
+    milvus_index, sample_chunks, sample_embeddings, embedding_dimension, mock_milvus_client
+):
+    """Test hybrid search with weighted reranker."""
+    mock_milvus_client.has_collection.return_value = True
+    await milvus_index.add_chunks(sample_chunks, sample_embeddings)
+
+    # Mock hybrid search results
+    mock_milvus_client.hybrid_search.return_value = [
+        [
+            {
+                "id": 0,
+                "distance": 0.1,
+                "entity": {"chunk_content": {"content": "mock chunk 1", "metadata": {"document_id": "doc1"}}},
+            },
+            {
+                "id": 1,
+                "distance": 0.2,
+                "entity": {"chunk_content": {"content": "mock chunk 2", "metadata": {"document_id": "doc2"}}},
+            },
+        ]
+    ]
+
+    # Test hybrid search with weighted reranker
+    query_embedding = np.random.rand(embedding_dimension).astype(np.float32)
+    query_string = "test query"
+    response = await milvus_index.query_hybrid(
+        embedding=query_embedding,
+        query_string=query_string,
+        k=2,
+        score_threshold=0.0,
+        reranker_type="weighted",
+        reranker_params={"alpha": 0.7},
+    )
+
+    assert isinstance(response, QueryChunksResponse)
+    assert len(response.chunks) == 2
+    assert len(response.scores) == 2
+
+    # Verify hybrid search was called with correct parameters
+    mock_milvus_client.hybrid_search.assert_called_once()
+    call_args = mock_milvus_client.hybrid_search.call_args
+    ranker = call_args[1]["ranker"]
+    assert ranker is not None
+
+
+async def test_query_hybrid_search_default_rrf(
+    milvus_index, sample_chunks, sample_embeddings, embedding_dimension, mock_milvus_client
+):
+    """Test hybrid search with default RRF reranker (no reranker_type specified)."""
+    mock_milvus_client.has_collection.return_value = True
+    await milvus_index.add_chunks(sample_chunks, sample_embeddings)
+
+    # Mock hybrid search results
+    mock_milvus_client.hybrid_search.return_value = [
+        [
+            {
+                "id": 0,
+                "distance": 0.1,
+                "entity": {"chunk_content": {"content": "mock chunk 1", "metadata": {"document_id": "doc1"}}},
+            },
+        ]
+    ]
+
+    # Test hybrid search with default reranker (should be RRF)
+    query_embedding = np.random.rand(embedding_dimension).astype(np.float32)
+    query_string = "test query"
+    response = await milvus_index.query_hybrid(
+        embedding=query_embedding,
+        query_string=query_string,
+        k=1,
+        score_threshold=0.0,
+        reranker_type="unknown_type",  # Should default to RRF
+        reranker_params=None,  # Should use default impact_factor
+    )
+
+    assert isinstance(response, QueryChunksResponse)
+    assert len(response.chunks) == 1
+
+    # Verify hybrid search was called with RRF reranker
+    mock_milvus_client.hybrid_search.assert_called_once()
+    call_args = mock_milvus_client.hybrid_search.call_args
+    ranker = call_args[1]["ranker"]
+    assert ranker is not None

tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py

@@ -12,7 +12,7 @@ import numpy as np
 import pytest
 
 from llama_stack.apis.vector_dbs import VectorDB
-from llama_stack.apis.vector_io import Chunk, QueryChunksResponse
+from llama_stack.apis.vector_io import Chunk, ChunkMetadata, QueryChunksResponse, VectorStoreContent
 from llama_stack.providers.remote.vector_io.milvus.milvus import VECTOR_DBS_PREFIX
 
 # This test is a unit test for the inline VectoerIO providers. This should only contain
@@ -294,3 +294,35 @@ async def test_delete_openai_vector_store_file_from_storage(vector_io_adapter, t
     assert loaded_file_info == {}
     loaded_contents = await vector_io_adapter._load_openai_vector_store_file_contents(store_id, file_id)
     assert loaded_contents == []
+
+
+async def test_chunk_to_vector_store_content_with_new_fields(vector_io_adapter):
+    sample_chunk_metadata = ChunkMetadata(
+        chunk_id="chunk123",
+        document_id="doc456",
+        source="test_source",
+        created_timestamp=1625133600,
+        updated_timestamp=1625133600,
+        chunk_window="0-100",
+        chunk_tokenizer="test_tokenizer",
+        chunk_embedding_model="dummy_model",
+        chunk_embedding_dimension=384,
+        content_token_count=100,
+        metadata_token_count=100,
+    )
+
+    sample_chunk = Chunk(
+        content="hello world", metadata={"lang": "en"}, embedding=[0.5, 0.7, 0.9], chunk_metadata=sample_chunk_metadata
+    )
+
+    vsc_list: VectorStoreContent = vector_io_adapter._chunk_to_vector_store_content(sample_chunk)
+    assert isinstance(vsc_list, list)
+    assert len(vsc_list) > 0
+
+    vsc = vsc_list[0]
+    assert vsc.text == "hello world"
+    assert vsc.type == "text"
+    assert vsc.metadata == {"lang": "en"}
+    assert vsc.chunk_metadata == sample_chunk_metadata
+    assert vsc.embedding == [0.5, 0.7, 0.9]
+    assert vsc.created_timestamp == 1625133600