mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-06-28 02:53:30 +00:00
chore: fixed formatting issues
This commit is contained in:
parent
f60c3c4acf
commit
4b32cfa846
1 changed files with 11 additions and 25 deletions
|
@ -14,12 +14,11 @@ import pytest_asyncio
|
||||||
from llama_stack.apis.inference import EmbeddingsResponse, Inference
|
from llama_stack.apis.inference import EmbeddingsResponse, Inference
|
||||||
from llama_stack.apis.vector_dbs import VectorDB
|
from llama_stack.apis.vector_dbs import VectorDB
|
||||||
from llama_stack.apis.vector_io import Chunk, QueryChunksResponse
|
from llama_stack.apis.vector_io import Chunk, QueryChunksResponse
|
||||||
|
from llama_stack.providers.inline.vector_io.faiss.config import FaissVectorIOConfig
|
||||||
from llama_stack.providers.inline.vector_io.faiss.faiss import (
|
from llama_stack.providers.inline.vector_io.faiss.faiss import (
|
||||||
FaissIndex,
|
FaissIndex,
|
||||||
FaissVectorIOAdapter,
|
FaissVectorIOAdapter,
|
||||||
)
|
)
|
||||||
from llama_stack.providers.inline.vector_io.faiss.config import FaissVectorIOConfig
|
|
||||||
|
|
||||||
# This test is a unit test for the FaissVectorIOAdapter class. This should only contain
|
# This test is a unit test for the FaissVectorIOAdapter class. This should only contain
|
||||||
# tests which are specific to this class. More general (API-level) tests should be placed in
|
# tests which are specific to this class. More general (API-level) tests should be placed in
|
||||||
|
@ -51,16 +50,8 @@ def vector_db_id():
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def sample_chunks():
|
def sample_chunks():
|
||||||
return [
|
return [
|
||||||
Chunk(
|
Chunk(content="MOCK text content 1", mime_type="text/plain", metadata={"document_id": "mock-doc-1"}),
|
||||||
content="MOCK text content 1",
|
Chunk(content="MOCK text content 1", mime_type="text/plain", metadata={"document_id": "mock-doc-2"}),
|
||||||
mime_type="text/plain",
|
|
||||||
metadata={"document_id": "mock-doc-1"}
|
|
||||||
),
|
|
||||||
Chunk(
|
|
||||||
content="MOCK text content 1",
|
|
||||||
mime_type="text/plain",
|
|
||||||
metadata={"document_id": "mock-doc-2"}
|
|
||||||
)
|
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
@ -107,28 +98,23 @@ async def faiss_adapter(faiss_config, mock_inference_api) -> FaissVectorIOAdapte
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_faiss_query_vector_returns_infinity_when_query_and_embedding_are_identical(faiss_index, sample_chunks, sample_embeddings, embedding_dimension):
|
async def test_faiss_query_vector_returns_infinity_when_query_and_embedding_are_identical(
|
||||||
|
faiss_index, sample_chunks, sample_embeddings, embedding_dimension
|
||||||
|
):
|
||||||
await faiss_index.add_chunks(sample_chunks, sample_embeddings)
|
await faiss_index.add_chunks(sample_chunks, sample_embeddings)
|
||||||
query_embedding = np.random.rand(embedding_dimension).astype(np.float32)
|
query_embedding = np.random.rand(embedding_dimension).astype(np.float32)
|
||||||
|
|
||||||
with patch.object(faiss_index.index, 'search') as mock_search:
|
with patch.object(faiss_index.index, "search") as mock_search:
|
||||||
mock_search.return_value = (
|
mock_search.return_value = (np.array([[0.0, 0.1]]), np.array([[0, 1]]))
|
||||||
np.array([[0.0, 0.1]]),
|
|
||||||
np.array([[0, 1]])
|
|
||||||
)
|
|
||||||
|
|
||||||
response = await faiss_index.query_vector(
|
response = await faiss_index.query_vector(embedding=query_embedding, k=2, score_threshold=0.0)
|
||||||
embedding=query_embedding,
|
|
||||||
k=2,
|
|
||||||
score_threshold=0.0
|
|
||||||
)
|
|
||||||
|
|
||||||
assert isinstance(response, QueryChunksResponse)
|
assert isinstance(response, QueryChunksResponse)
|
||||||
assert len(response.chunks) == 2
|
assert len(response.chunks) == 2
|
||||||
assert len(response.scores) == 2
|
assert len(response.scores) == 2
|
||||||
|
|
||||||
assert response.scores[0] == float("inf") # infinity (1.0 / 0.0)
|
assert response.scores[0] == float("inf") # infinity (1.0 / 0.0)
|
||||||
assert response.scores[1] == 10.0 # (1.0 / 0.1 = 10.0)
|
assert response.scores[1] == 10.0 # (1.0 / 0.1 = 10.0)
|
||||||
|
|
||||||
assert response.chunks[0] == sample_chunks[0]
|
assert response.chunks[0] == sample_chunks[0]
|
||||||
assert response.chunks[1] == sample_chunks[1]
|
assert response.chunks[1] == sample_chunks[1]
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue