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,