feat: update search for vector_stores (#2441)

Updated the `search` functionality return response to match openai. 

## Test Plan
```
pytest -sv --stack-config=http://localhost:8321 tests/integration/vector_io/test_openai_vector_stores.py --embedding-model all-MiniLM-L6-v2
```

docs/_static/llama-stack-spec.html

@@ -3864,7 +3864,7 @@
                         "content": {
                             "application/json": {
                                 "schema": {
-                                    "$ref": "#/components/schemas/VectorStoreSearchResponse"
+                                    "$ref": "#/components/schemas/VectorStoreSearchResponsePage"
                                 }
                             }
                         }
@@ -13132,7 +13132,70 @@
                 ],
                 "title": "OpenaiSearchVectorStoreRequest"
             },
+            "VectorStoreContent": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "text"
+                    },
+                    "text": {
+                        "type": "string"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type",
+                    "text"
+                ],
+                "title": "VectorStoreContent"
+            },
             "VectorStoreSearchResponse": {
+                "type": "object",
+                "properties": {
+                    "file_id": {
+                        "type": "string"
+                    },
+                    "filename": {
+                        "type": "string"
+                    },
+                    "score": {
+                        "type": "number"
+                    },
+                    "attributes": {
+                        "type": "object",
+                        "additionalProperties": {
+                            "oneOf": [
+                                {
+                                    "type": "string"
+                                },
+                                {
+                                    "type": "number"
+                                },
+                                {
+                                    "type": "boolean"
+                                }
+                            ]
+                        }
+                    },
+                    "content": {
+                        "type": "array",
+                        "items": {
+                            "$ref": "#/components/schemas/VectorStoreContent"
+                        }
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "file_id",
+                    "filename",
+                    "score",
+                    "content"
+                ],
+                "title": "VectorStoreSearchResponse",
+                "description": "Response from searching a vector store."
+            },
+            "VectorStoreSearchResponsePage": {
                 "type": "object",
                 "properties": {
                     "object": {
@@ -13145,29 +13208,7 @@
                     "data": {
                         "type": "array",
                         "items": {
-                            "type": "object",
-                            "additionalProperties": {
-                                "oneOf": [
-                                    {
-                                        "type": "null"
-                                    },
-                                    {
-                                        "type": "boolean"
-                                    },
-                                    {
-                                        "type": "number"
-                                    },
-                                    {
-                                        "type": "string"
-                                    },
-                                    {
-                                        "type": "array"
-                                    },
-                                    {
-                                        "type": "object"
-                                    }
-                                ]
-                            }
+                            "$ref": "#/components/schemas/VectorStoreSearchResponse"
                         }
                     },
                     "has_more": {
@@ -13185,7 +13226,7 @@
                     "data",
                     "has_more"
                 ],
-                "title": "VectorStoreSearchResponse",
+                "title": "VectorStoreSearchResponsePage",
                 "description": "Response from searching a vector store."
             },
             "OpenaiUpdateVectorStoreRequest": {

docs/_static/llama-stack-spec.yaml

@@ -2734,7 +2734,7 @@ paths:
           content:
             application/json:
               schema:
-                $ref: '#/components/schemas/VectorStoreSearchResponse'
+                $ref: '#/components/schemas/VectorStoreSearchResponsePage'
         '400':
           $ref: '#/components/responses/BadRequest400'
         '429':
@@ -9190,7 +9190,48 @@ components:
       required:
         - query
       title: OpenaiSearchVectorStoreRequest
+    VectorStoreContent:
+      type: object
+      properties:
+        type:
+          type: string
+          const: text
+        text:
+          type: string
+      additionalProperties: false
+      required:
+        - type
+        - text
+      title: VectorStoreContent
     VectorStoreSearchResponse:
+      type: object
+      properties:
+        file_id:
+          type: string
+        filename:
+          type: string
+        score:
+          type: number
+        attributes:
+          type: object
+          additionalProperties:
+            oneOf:
+              - type: string
+              - type: number
+              - type: boolean
+        content:
+          type: array
+          items:
+            $ref: '#/components/schemas/VectorStoreContent'
+      additionalProperties: false
+      required:
+        - file_id
+        - filename
+        - score
+        - content
+      title: VectorStoreSearchResponse
+      description: Response from searching a vector store.
+    VectorStoreSearchResponsePage:
       type: object
       properties:
         object:
@@ -9201,15 +9242,7 @@ components:
         data:
           type: array
           items:
-            type: object
-            additionalProperties:
-              oneOf:
-                - type: 'null'
-                - type: boolean
-                - type: number
-                - type: string
-                - type: array
-                - type: object
+            $ref: '#/components/schemas/VectorStoreSearchResponse'
         has_more:
           type: boolean
           default: false
@@ -9221,7 +9254,7 @@ components:
         - search_query
         - data
         - has_more
-      title: VectorStoreSearchResponse
+      title: VectorStoreSearchResponsePage
       description: Response from searching a vector store.
     OpenaiUpdateVectorStoreRequest:
       type: object

llama_stack/apis/vector_io/vector_io.py

@@ -8,7 +8,7 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
-from typing import Any, Protocol, runtime_checkable
+from typing import Any, Literal, Protocol, runtime_checkable
 
 from pydantic import BaseModel, Field
 
@@ -96,13 +96,30 @@ class VectorStoreSearchRequest(BaseModel):
     rewrite_query: bool = False
 
 
+@json_schema_type
+class VectorStoreContent(BaseModel):
+    type: Literal["text"]
+    text: str
+
+
 @json_schema_type
 class VectorStoreSearchResponse(BaseModel):
     """Response from searching a vector store."""
 
+    file_id: str
+    filename: str
+    score: float
+    attributes: dict[str, str | float | bool] | None = None
+    content: list[VectorStoreContent]
+
+
+@json_schema_type
+class VectorStoreSearchResponsePage(BaseModel):
+    """Response from searching a vector store."""
+
     object: str = "vector_store.search_results.page"
     search_query: str
-    data: list[dict[str, Any]]
+    data: list[VectorStoreSearchResponse]
     has_more: bool = False
     next_page: str | None = None
 
@@ -259,7 +276,7 @@ class VectorIO(Protocol):
         max_num_results: int | None = 10,
         ranking_options: dict[str, Any] | None = None,
         rewrite_query: bool | None = False,
-    ) -> VectorStoreSearchResponse:
+    ) -> VectorStoreSearchResponsePage:
         """Search for chunks in a vector store.
 
         Searches a vector store for relevant chunks based on a query and optional file attribute filters.

llama_stack/distribution/routers/vector_io.py

@@ -17,7 +17,7 @@ from llama_stack.apis.vector_io import (
     VectorStoreDeleteResponse,
     VectorStoreListResponse,
     VectorStoreObject,
-    VectorStoreSearchResponse,
+    VectorStoreSearchResponsePage,
 )
 from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import RoutingTable
@@ -242,7 +242,7 @@ class VectorIORouter(VectorIO):
         max_num_results: int | None = 10,
         ranking_options: dict[str, Any] | None = None,
         rewrite_query: bool | None = False,
-    ) -> VectorStoreSearchResponse:
+    ) -> VectorStoreSearchResponsePage:
         logger.debug(f"VectorIORouter.openai_search_vector_store: {vector_store_id}")
         # Route based on vector store ID
         provider = self.routing_table.get_provider_impl(vector_store_id)

llama_stack/providers/remote/vector_io/chroma/chroma.py

@@ -21,7 +21,7 @@ from llama_stack.apis.vector_io import (
     VectorStoreDeleteResponse,
     VectorStoreListResponse,
     VectorStoreObject,
-    VectorStoreSearchResponse,
+    VectorStoreSearchResponsePage,
 )
 from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
 from llama_stack.providers.inline.vector_io.chroma import ChromaVectorIOConfig as InlineChromaVectorIOConfig
@@ -239,5 +239,5 @@ class ChromaVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
         max_num_results: int | None = 10,
         ranking_options: dict[str, Any] | None = None,
         rewrite_query: bool | None = False,
-    ) -> VectorStoreSearchResponse:
+    ) -> VectorStoreSearchResponsePage:
         raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")

llama_stack/providers/remote/vector_io/milvus/milvus.py

@@ -23,7 +23,7 @@ from llama_stack.apis.vector_io import (
     VectorStoreDeleteResponse,
     VectorStoreListResponse,
     VectorStoreObject,
-    VectorStoreSearchResponse,
+    VectorStoreSearchResponsePage,
 )
 from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
 from llama_stack.providers.inline.vector_io.milvus import MilvusVectorIOConfig as InlineMilvusVectorIOConfig
@@ -237,7 +237,7 @@ class MilvusVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
         max_num_results: int | None = 10,
         ranking_options: dict[str, Any] | None = None,
         rewrite_query: bool | None = False,
-    ) -> VectorStoreSearchResponse:
+    ) -> VectorStoreSearchResponsePage:
         raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
 
 

llama_stack/providers/remote/vector_io/qdrant/qdrant.py

@@ -21,7 +21,7 @@ from llama_stack.apis.vector_io import (
     VectorStoreDeleteResponse,
     VectorStoreListResponse,
     VectorStoreObject,
-    VectorStoreSearchResponse,
+    VectorStoreSearchResponsePage,
 )
 from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
 from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig as InlineQdrantVectorIOConfig
@@ -239,5 +239,5 @@ class QdrantVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
         max_num_results: int | None = 10,
         ranking_options: dict[str, Any] | None = None,
         rewrite_query: bool | None = False,
-    ) -> VectorStoreSearchResponse:
+    ) -> VectorStoreSearchResponsePage:
         raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")

llama_stack/providers/utils/memory/openai_vector_store_mixin.py

@@ -13,10 +13,12 @@ from typing import Any
 from llama_stack.apis.vector_dbs import VectorDB
 from llama_stack.apis.vector_io import (
     QueryChunksResponse,
+    VectorStoreContent,
     VectorStoreDeleteResponse,
     VectorStoreListResponse,
     VectorStoreObject,
     VectorStoreSearchResponse,
+    VectorStoreSearchResponsePage,
 )
 
 logger = logging.getLogger(__name__)
@@ -85,7 +87,6 @@ class OpenAIVectorStoreMixin(ABC):
         provider_vector_db_id: str | None = None,
     ) -> VectorStoreObject:
         """Creates a vector store."""
-        print("IN OPENAI VECTOR STORE MIXIN, openai_create_vector_store")
         # store and vector_db have the same id
         store_id = name or str(uuid.uuid4())
         created_at = int(time.time())
@@ -281,7 +282,7 @@ class OpenAIVectorStoreMixin(ABC):
         ranking_options: dict[str, Any] | None = None,
         rewrite_query: bool | None = False,
         # search_mode: Literal["keyword", "vector", "hybrid"] = "vector",
-    ) -> VectorStoreSearchResponse:
+    ) -> VectorStoreSearchResponsePage:
         """Search for chunks in a vector store."""
         # TODO: Add support in the API for this
         search_mode = "vector"
@@ -312,7 +313,7 @@ class OpenAIVectorStoreMixin(ABC):
 
             # Convert response to OpenAI format
             data = []
-            for i, (chunk, score) in enumerate(zip(response.chunks, response.scores, strict=False)):
+            for chunk, score in zip(response.chunks, response.scores, strict=False):
                 # Apply score based filtering
                 if score < score_threshold:
                     continue
@@ -323,18 +324,46 @@ class OpenAIVectorStoreMixin(ABC):
                     if not self._matches_filters(chunk.metadata, filters):
                         continue
 
-                chunk_data = {
-                    "id": f"chunk_{i}",
-                    "object": "vector_store.search_result",
-                    "score": score,
-                    "content": chunk.content.content if hasattr(chunk.content, "content") else str(chunk.content),
-                    "metadata": chunk.metadata,
-                }
-                data.append(chunk_data)
+                # content is InterleavedContent
+                if isinstance(chunk.content, str):
+                    content = [
+                        VectorStoreContent(
+                            type="text",
+                            text=chunk.content,
+                        )
+                    ]
+                elif isinstance(chunk.content, list):
+                    # TODO: Add support for other types of content
+                    content = [
+                        VectorStoreContent(
+                            type="text",
+                            text=item.text,
+                        )
+                        for item in chunk.content
+                        if item.type == "text"
+                    ]
+                else:
+                    if chunk.content.type != "text":
+                        raise ValueError(f"Unsupported content type: {chunk.content.type}")
+                    content = [
+                        VectorStoreContent(
+                            type="text",
+                            text=chunk.content.text,
+                        )
+                    ]
+
+                response_data_item = VectorStoreSearchResponse(
+                    file_id=chunk.metadata.get("file_id", ""),
+                    filename=chunk.metadata.get("filename", ""),
+                    score=score,
+                    attributes=chunk.metadata,
+                    content=content,
+                )
+                data.append(response_data_item)
                 if len(data) >= max_num_results:
                     break
 
-            return VectorStoreSearchResponse(
+            return VectorStoreSearchResponsePage(
                 search_query=search_query,
                 data=data,
                 has_more=False,  # For simplicity, we don't implement pagination here
@@ -344,7 +373,7 @@ class OpenAIVectorStoreMixin(ABC):
         except Exception as e:
             logger.error(f"Error searching vector store {vector_store_id}: {e}")
             # Return empty results on error
-            return VectorStoreSearchResponse(
+            return VectorStoreSearchResponsePage(
                 search_query=search_query,
                 data=[],
                 has_more=False,

tests/integration/vector_io/test_openai_vector_stores.py

@@ -34,6 +34,13 @@ def openai_client(client_with_models):
     return OpenAI(base_url=base_url, api_key="fake")
 
 
+@pytest.fixture(params=["openai_client"])  # , "llama_stack_client"])
+def compat_client(request, client_with_models):
+    if request.param == "openai_client" and isinstance(client_with_models, LlamaStackAsLibraryClient):
+        pytest.skip("OpenAI client tests not supported with library client")
+    return request.getfixturevalue(request.param)
+
+
 @pytest.fixture(scope="session")
 def sample_chunks():
     return [
@@ -57,29 +64,29 @@ def sample_chunks():
 
 
 @pytest.fixture(scope="function")
-def openai_client_with_empty_stores(openai_client):
+def compat_client_with_empty_stores(compat_client):
     def clear_vector_stores():
         # List and delete all existing vector stores
         try:
-            response = openai_client.vector_stores.list()
+            response = compat_client.vector_stores.list()
             for store in response.data:
-                openai_client.vector_stores.delete(vector_store_id=store.id)
+                compat_client.vector_stores.delete(vector_store_id=store.id)
         except Exception:
             # If the API is not available or fails, just continue
             logger.warning("Failed to clear vector stores")
             pass
 
     clear_vector_stores()
-    yield openai_client
+    yield compat_client
 
     # Clean up after the test
     clear_vector_stores()
 
 
-def test_openai_create_vector_store(openai_client_with_empty_stores, client_with_models):
+def test_openai_create_vector_store(compat_client_with_empty_stores, client_with_models):
     """Test creating a vector store using OpenAI API."""
     skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
-    client = openai_client_with_empty_stores
+    client = compat_client_with_empty_stores
 
     # Create a vector store
     vector_store = client.vector_stores.create(
@@ -96,11 +103,11 @@ def test_openai_create_vector_store(openai_client_with_empty_stores, client_with
     assert hasattr(vector_store, "created_at")
 
 
-def test_openai_list_vector_stores(openai_client_with_empty_stores, client_with_models):
+def test_openai_list_vector_stores(compat_client_with_empty_stores, client_with_models):
     """Test listing vector stores using OpenAI API."""
     skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
 
-    client = openai_client_with_empty_stores
+    client = compat_client_with_empty_stores
 
     # Create a few vector stores
     store1 = client.vector_stores.create(name="store1", metadata={"type": "test"})
@@ -123,11 +130,11 @@ def test_openai_list_vector_stores(openai_client_with_empty_stores, client_with_
     assert len(limited_response.data) == 1
 
 
-def test_openai_retrieve_vector_store(openai_client_with_empty_stores, client_with_models):
+def test_openai_retrieve_vector_store(compat_client_with_empty_stores, client_with_models):
     """Test retrieving a specific vector store using OpenAI API."""
     skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
 
-    client = openai_client_with_empty_stores
+    client = compat_client_with_empty_stores
 
     # Create a vector store
     created_store = client.vector_stores.create(name="retrieve_test_store", metadata={"purpose": "retrieval_test"})
@@ -142,11 +149,11 @@ def test_openai_retrieve_vector_store(openai_client_with_empty_stores, client_wi
     assert retrieved_store.object == "vector_store"
 
 
-def test_openai_update_vector_store(openai_client_with_empty_stores, client_with_models):
+def test_openai_update_vector_store(compat_client_with_empty_stores, client_with_models):
     """Test modifying a vector store using OpenAI API."""
     skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
 
-    client = openai_client_with_empty_stores
+    client = compat_client_with_empty_stores
 
     # Create a vector store
     created_store = client.vector_stores.create(name="original_name", metadata={"version": "1.0"})
@@ -165,11 +172,11 @@ def test_openai_update_vector_store(openai_client_with_empty_stores, client_with
     assert modified_store.last_active_at > created_store.last_active_at
 
 
-def test_openai_delete_vector_store(openai_client_with_empty_stores, client_with_models):
+def test_openai_delete_vector_store(compat_client_with_empty_stores, client_with_models):
     """Test deleting a vector store using OpenAI API."""
     skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
 
-    client = openai_client_with_empty_stores
+    client = compat_client_with_empty_stores
 
     # Create a vector store
     created_store = client.vector_stores.create(name="delete_test_store", metadata={"purpose": "deletion_test"})
@@ -187,11 +194,11 @@ def test_openai_delete_vector_store(openai_client_with_empty_stores, client_with
         client.vector_stores.retrieve(vector_store_id=created_store.id)
 
 
-def test_openai_vector_store_search_empty(openai_client_with_empty_stores, client_with_models):
+def test_openai_vector_store_search_empty(compat_client_with_empty_stores, client_with_models):
     """Test searching an empty vector store using OpenAI API."""
     skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
 
-    client = openai_client_with_empty_stores
+    client = compat_client_with_empty_stores
 
     # Create a vector store
     vector_store = client.vector_stores.create(name="search_test_store", metadata={"purpose": "search_testing"})
@@ -208,15 +215,15 @@ def test_openai_vector_store_search_empty(openai_client_with_empty_stores, clien
     assert search_response.has_more is False
 
 
-def test_openai_vector_store_with_chunks(openai_client_with_empty_stores, client_with_models, sample_chunks):
+def test_openai_vector_store_with_chunks(compat_client_with_empty_stores, client_with_models, sample_chunks):
     """Test vector store functionality with actual chunks using both OpenAI and native APIs."""
     skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
 
-    openai_client = openai_client_with_empty_stores
+    compat_client = compat_client_with_empty_stores
     llama_client = client_with_models
 
     # Create a vector store using OpenAI API
-    vector_store = openai_client.vector_stores.create(name="chunks_test_store", metadata={"purpose": "chunks_testing"})
+    vector_store = compat_client.vector_stores.create(name="chunks_test_store", metadata={"purpose": "chunks_testing"})
 
     # Insert chunks using the native LlamaStack API (since OpenAI API doesn't have direct chunk insertion)
     llama_client.vector_io.insert(
@@ -225,7 +232,7 @@ def test_openai_vector_store_with_chunks(openai_client_with_empty_stores, client
     )
 
     # Search using OpenAI API
-    search_response = openai_client.vector_stores.search(
+    search_response = compat_client.vector_stores.search(
         vector_store_id=vector_store.id, query="What is Python programming language?", max_num_results=3
     )
     assert search_response is not None
@@ -233,18 +240,19 @@ def test_openai_vector_store_with_chunks(openai_client_with_empty_stores, client
 
     # The top result should be about Python (doc1)
     top_result = search_response.data[0]
-    assert "python" in top_result.content.lower() or "programming" in top_result.content.lower()
-    assert top_result.metadata["document_id"] == "doc1"
+    top_content = top_result.content[0].text
+    assert "python" in top_content.lower() or "programming" in top_content.lower()
+    assert top_result.attributes["document_id"] == "doc1"
 
     # Test filtering by metadata
-    filtered_search = openai_client.vector_stores.search(
+    filtered_search = compat_client.vector_stores.search(
         vector_store_id=vector_store.id, query="artificial intelligence", filters={"topic": "ai"}, max_num_results=5
     )
 
     assert filtered_search is not None
     # All results should have topic "ai"
     for result in filtered_search.data:
-        assert result.metadata["topic"] == "ai"
+        assert result.attributes["topic"] == "ai"
 
 
 @pytest.mark.parametrize(
@@ -257,18 +265,18 @@ def test_openai_vector_store_with_chunks(openai_client_with_empty_stores, client
     ],
 )
 def test_openai_vector_store_search_relevance(
-    openai_client_with_empty_stores, client_with_models, sample_chunks, test_case
+    compat_client_with_empty_stores, client_with_models, sample_chunks, test_case
 ):
     """Test that OpenAI vector store search returns relevant results for different queries."""
     skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
 
-    openai_client = openai_client_with_empty_stores
+    compat_client = compat_client_with_empty_stores
     llama_client = client_with_models
 
     query, expected_doc_id, expected_topic = test_case
 
     # Create a vector store
-    vector_store = openai_client.vector_stores.create(
+    vector_store = compat_client.vector_stores.create(
         name=f"relevance_test_{expected_doc_id}", metadata={"purpose": "relevance_testing"}
     )
 
@@ -279,7 +287,7 @@ def test_openai_vector_store_search_relevance(
     )
 
     # Search using OpenAI API
-    search_response = openai_client.vector_stores.search(
+    search_response = compat_client.vector_stores.search(
         vector_store_id=vector_store.id, query=query, max_num_results=4
     )
 
@@ -288,8 +296,9 @@ def test_openai_vector_store_search_relevance(
 
     # The top result should match the expected document
     top_result = search_response.data[0]
-    assert top_result.metadata["document_id"] == expected_doc_id
-    assert top_result.metadata["topic"] == expected_topic
+
+    assert top_result.attributes["document_id"] == expected_doc_id
+    assert top_result.attributes["topic"] == expected_topic
 
     # Verify score is included and reasonable
     assert isinstance(top_result.score, int | float)
@@ -297,16 +306,16 @@ def test_openai_vector_store_search_relevance(
 
 
 def test_openai_vector_store_search_with_ranking_options(
-    openai_client_with_empty_stores, client_with_models, sample_chunks
+    compat_client_with_empty_stores, client_with_models, sample_chunks
 ):
     """Test OpenAI vector store search with ranking options."""
     skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
 
-    openai_client = openai_client_with_empty_stores
+    compat_client = compat_client_with_empty_stores
     llama_client = client_with_models
 
     # Create a vector store
-    vector_store = openai_client.vector_stores.create(
+    vector_store = compat_client.vector_stores.create(
         name="ranking_test_store", metadata={"purpose": "ranking_testing"}
     )
 
@@ -318,7 +327,7 @@ def test_openai_vector_store_search_with_ranking_options(
 
     # Search with ranking options
     threshold = 0.1
-    search_response = openai_client.vector_stores.search(
+    search_response = compat_client.vector_stores.search(
         vector_store_id=vector_store.id,
         query="machine learning and artificial intelligence",
         max_num_results=3,
@@ -334,16 +343,16 @@ def test_openai_vector_store_search_with_ranking_options(
 
 
 def test_openai_vector_store_search_with_high_score_filter(
-    openai_client_with_empty_stores, client_with_models, sample_chunks
+    compat_client_with_empty_stores, client_with_models, sample_chunks
 ):
     """Test that searching with text very similar to a document and high score threshold returns only that document."""
     skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
 
-    openai_client = openai_client_with_empty_stores
+    compat_client = compat_client_with_empty_stores
     llama_client = client_with_models
 
     # Create a vector store
-    vector_store = openai_client.vector_stores.create(
+    vector_store = compat_client.vector_stores.create(
         name="high_score_filter_test", metadata={"purpose": "high_score_filtering"}
     )
 
@@ -358,7 +367,7 @@ def test_openai_vector_store_search_with_high_score_filter(
     query = "Python is a high-level programming language with code readability and fewer lines than C++ or Java"
 
     # picking up thrshold to be slightly higher than the second result
-    search_response = openai_client.vector_stores.search(
+    search_response = compat_client.vector_stores.search(
         vector_store_id=vector_store.id,
         query=query,
         max_num_results=3,
@@ -367,7 +376,7 @@ def test_openai_vector_store_search_with_high_score_filter(
     threshold = search_response.data[1].score + 0.0001
 
     # we expect only one result with the requested threshold
-    search_response = openai_client.vector_stores.search(
+    search_response = compat_client.vector_stores.search(
         vector_store_id=vector_store.id,
         query=query,
         max_num_results=10,  # Allow more results but expect filtering
@@ -379,25 +388,26 @@ def test_openai_vector_store_search_with_high_score_filter(
 
     # The top result should be the Python document (doc1)
     top_result = search_response.data[0]
-    assert top_result.metadata["document_id"] == "doc1"
-    assert top_result.metadata["topic"] == "programming"
+    assert top_result.attributes["document_id"] == "doc1"
+    assert top_result.attributes["topic"] == "programming"
     assert top_result.score >= threshold
 
     # Verify the content contains Python-related terms
-    assert "python" in top_result.content.lower() or "programming" in top_result.content.lower()
+    top_content = top_result.content[0].text
+    assert "python" in top_content.lower() or "programming" in top_content.lower()
 
 
 def test_openai_vector_store_search_with_max_num_results(
-    openai_client_with_empty_stores, client_with_models, sample_chunks
+    compat_client_with_empty_stores, client_with_models, sample_chunks
 ):
     """Test OpenAI vector store search with max_num_results."""
     skip_if_provider_doesnt_support_openai_vector_stores(client_with_models)
 
-    openai_client = openai_client_with_empty_stores
+    compat_client = compat_client_with_empty_stores
     llama_client = client_with_models
 
     # Create a vector store
-    vector_store = openai_client.vector_stores.create(
+    vector_store = compat_client.vector_stores.create(
         name="max_num_results_test_store", metadata={"purpose": "max_num_results_testing"}
     )
 
@@ -408,7 +418,7 @@ def test_openai_vector_store_search_with_max_num_results(
     )
 
     # Search with max_num_results
-    search_response = openai_client.vector_stores.search(
+    search_response = compat_client.vector_stores.search(
         vector_store_id=vector_store.id,
         query="machine learning and artificial intelligence",
         max_num_results=2,