diff --git a/docs/_static/llama-stack-spec.html b/docs/_static/llama-stack-spec.html index a1a3217c4..96de04ec9 100644 --- a/docs/_static/llama-stack-spec.html +++ b/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": { diff --git a/docs/_static/llama-stack-spec.yaml b/docs/_static/llama-stack-spec.yaml index 15593d060..b2fe870be 100644 --- a/docs/_static/llama-stack-spec.yaml +++ b/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 diff --git a/llama_stack/apis/vector_io/vector_io.py b/llama_stack/apis/vector_io/vector_io.py index c14a88c5e..1c8ae4dab 100644 --- a/llama_stack/apis/vector_io/vector_io.py +++ b/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. diff --git a/llama_stack/distribution/routers/vector_io.py b/llama_stack/distribution/routers/vector_io.py index 601109963..3d65aef24 100644 --- a/llama_stack/distribution/routers/vector_io.py +++ b/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) diff --git a/llama_stack/providers/remote/vector_io/chroma/chroma.py b/llama_stack/providers/remote/vector_io/chroma/chroma.py index 5f5be539d..0d8451eb2 100644 --- a/llama_stack/providers/remote/vector_io/chroma/chroma.py +++ b/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") diff --git a/llama_stack/providers/remote/vector_io/milvus/milvus.py b/llama_stack/providers/remote/vector_io/milvus/milvus.py index ae59af599..8ae74aedc 100644 --- a/llama_stack/providers/remote/vector_io/milvus/milvus.py +++ b/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") diff --git a/llama_stack/providers/remote/vector_io/qdrant/qdrant.py b/llama_stack/providers/remote/vector_io/qdrant/qdrant.py index 9e802fd6a..10f3b5b0d 100644 --- a/llama_stack/providers/remote/vector_io/qdrant/qdrant.py +++ b/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") diff --git a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py index c2a63f3c5..398075f57 100644 --- a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py +++ b/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, diff --git a/tests/integration/vector_io/test_openai_vector_stores.py b/tests/integration/vector_io/test_openai_vector_stores.py index d67c35e69..a67582a07 100644 --- a/tests/integration/vector_io/test_openai_vector_stores.py +++ b/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,