phoenix-oss/llama-stack
1
0
Fork 0
forked from phoenix-oss/llama-stack-mirror
Code Pull requests Releases Packages 2 Activity Actions

chore: Get sqlite_vec and vector_store unit tests passing (#1413)

Browse source
This commit is contained in:
Ben Browning 2025-03-05 13:20:13 -05:00 • committed by GitHub
parent 77d323c2f8
commit 00570fde31
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 7 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

3
tests/unit/providers/vector_io/test_sqlite_vec.py
View file

@ -9,6 +9,7 @@ import sqlite3
import numpy as np import numpy as np
import pytest import pytest
import pytest_asyncio
import sqlite_vec import sqlite_vec
from llama_stack.apis.vector_io import Chunk, QueryChunksResponse from llama_stack.apis.vector_io import Chunk, QueryChunksResponse
@ -48,7 +49,7 @@ def sqlite_connection(loop):
conn.close() conn.close()
@pytest.fixture(scope="session", autouse=True) @pytest_asyncio.fixture(scope="session", autouse=True)
async def sqlite_vec_index(sqlite_connection): async def sqlite_vec_index(sqlite_connection):
return await SQLiteVecIndex.create(dimension=EMBEDDING_DIMENSION, connection=sqlite_connection, bank_id="test_bank") return await SQLiteVecIndex.create(dimension=EMBEDDING_DIMENSION, connection=sqlite_connection, bank_id="test_bank")

8
tests/unit/rag/test_vector_store.py
View file

@ -15,6 +15,8 @@ from llama_stack.apis.tools import RAGDocument
from llama_stack.providers.utils.memory.vector_store import URL, content_from_doc from llama_stack.providers.utils.memory.vector_store import URL, content_from_doc
DUMMY_PDF_PATH = Path(os.path.abspath(__file__)).parent / "fixtures" / "dummy.pdf" DUMMY_PDF_PATH = Path(os.path.abspath(__file__)).parent / "fixtures" / "dummy.pdf"
# Depending on the machine, this can get parsed a couple of ways
DUMMY_PDF_TEXT_CHOICES = ["Dummy PDF file", "Dumm y PDF file"]
def read_file(file_path: str) -> bytes: def read_file(file_path: str) -> bytes:
@ -45,7 +47,7 @@ class TestVectorStore:
metadata={}, metadata={},
) )
content = await content_from_doc(doc) content = await content_from_doc(doc)
assert content == "Dumm y PDF file" assert content in DUMMY_PDF_TEXT_CHOICES
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_downloads_pdf_and_returns_content(self): async def test_downloads_pdf_and_returns_content(self):
@ -58,7 +60,7 @@ class TestVectorStore:
metadata={}, metadata={},
) )
content = await content_from_doc(doc) content = await content_from_doc(doc)
assert content == "Dumm y PDF file" assert content in DUMMY_PDF_TEXT_CHOICES
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_downloads_pdf_and_returns_content_with_url_object(self): async def test_downloads_pdf_and_returns_content_with_url_object(self):
@ -73,4 +75,4 @@ class TestVectorStore:
metadata={}, metadata={},
) )
content = await content_from_doc(doc) content = await content_from_doc(doc)
assert content == "Dumm y PDF file" assert content in DUMMY_PDF_TEXT_CHOICES