mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-06-27 18:50:41 +00:00
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 ```
This commit is contained in:
Hardik Shah
GitHub
parent
35c2817d0a
commit
0bc1747ed8
9 changed files with 236 additions and 106 deletions
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue