From 49f831e87053f1d36420f846a1b4bf7a273de4c7 Mon Sep 17 00:00:00 2001 From: Laurent Sorber Date: Thu, 15 Aug 2024 16:14:26 +0200 Subject: [PATCH] test: improve test speed --- README.md | 4 +--- pyproject.toml | 52 ++------------------------------------------- tests/conftest.py | 34 +++++++++++++++++++++++++++++ tests/test_basic.py | 12 +++++------ 4 files changed, 42 insertions(+), 60 deletions(-) create mode 100644 tests/conftest.py diff --git a/README.md b/README.md index dee1cef..1c96252 100644 --- a/README.md +++ b/README.md @@ -2,9 +2,7 @@ # RAGLite -A Retrieval-Augmented Generation (RAG) library for SQLite. - -⚠️ This project is a work in progress! +RAGLite is a Python package for Retrieval-Augmented Generation (RAG) with SQLite. ## Features diff --git a/pyproject.toml b/pyproject.toml index a694fde..b4e3c84 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api" [tool.poetry] # https://python-poetry.org/docs/pyproject/ name = "raglite" version = "0.0.0" -description = "A RAG extension for SQLite." +description = "A Python package for Retrieval-Augmented Generation (RAG) with SQLite." authors = ["Laurent Sorber "] readme = "README.md" repository = "https://github.com/radix-ai/raglite" @@ -123,55 +123,7 @@ src = ["src", "tests"] target-version = "py310" [tool.ruff.lint] -select = [ - "A", - "ASYNC", - "B", - "BLE", - "C4", - "C90", - "D", - "DTZ", - "E", - "EM", - "ERA", - "F", - "FBT", - "FLY", - "FURB", - "G", - "I", - "ICN", - "INP", - "INT", - "ISC", - "LOG", - "N", - "NPY", - "PERF", - "PGH", - "PIE", - "PL", - "PT", - "PTH", - "PYI", - "Q", - "RET", - "RSE", - "RUF", - "S", - "SIM", - "SLF", - "SLOT", - "T10", - "T20", - "TCH", - "TID", - "TRY", - "UP", - "W", - "YTT", -] +select = ["A", "ASYNC", "B", "BLE", "C4", "C90", "D", "DTZ", "E", "EM", "ERA", "F", "FBT", "FLY", "FURB", "G", "I", "ICN", "INP", "INT", "ISC", "LOG", "N", "NPY", "PERF", "PGH", "PIE", "PL", "PT", "PTH", "PYI", "Q", "RET", "RSE", "RUF", "S", "SIM", "SLF", "SLOT", "T10", "T20", "TCH", "TID", "TRY", "UP", "W", "YTT"] ignore = ["D203", "D213", "E501", "RET504", "RUF002", "S101", "S307"] unfixable = ["ERA001", "F401", "F841", "T201", "T203"] diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..52a5d6d --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,34 @@ +"""Fixtures for the tests.""" + +import pytest +from llama_cpp import Llama, LlamaRAMCache # type: ignore[attr-defined] + +from raglite import RAGLiteConfig + + +@pytest.fixture() +def simple_config() -> RAGLiteConfig: + """Create a lightweight in-memory config for testing.""" + # Use a lightweight LLM. + llm = Llama.from_pretrained( + repo_id="bartowski/Phi-3.1-mini-4k-instruct-GGUF", # https://huggingface.co/microsoft/Phi-3-mini-4k-instruct + filename="*Q4_K_M.gguf", + n_ctx=4096, # 0 = Use the model's context size (default is 512). + n_gpu_layers=-1, # -1 = Offload all layers to the GPU (default is 0). + verbose=False, + ) + llm.set_cache(LlamaRAMCache()) + # Use a lightweight embedder. + embedder = Llama.from_pretrained( + repo_id="ChristianAzinn/snowflake-arctic-embed-xs-gguf", # https://github.com/Snowflake-Labs/arctic-embed + filename="*f16.GGUF", + n_ctx=0, # 0 = Use the model's context size (default is 512). + n_gpu_layers=-1, # -1 = Offload all layers to the GPU (default is 0). + verbose=False, + embedding=True, + ) + # Use an in-memory SQLite database. + db_url = "sqlite:///:memory:" + # Create the config. + config = RAGLiteConfig(llm=llm, embedder=embedder, db_url=db_url) + return config diff --git a/tests/test_basic.py b/tests/test_basic.py index fb5dfdf..f640e3c 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -11,22 +11,20 @@ ) -def test_insert_index_search() -> None: +def test_insert_index_search(simple_config: RAGLiteConfig) -> None: """Test inserting a document, updating the vector index, and searching for a query.""" - # Run this test with an in-memory SQLite database. - in_memory_db = RAGLiteConfig(db_url="sqlite:///:memory:") # Insert a document. doc_path = Path(__file__).parent / "specrel.pdf" # Einstein's special relativity paper. - insert_document(doc_path, config=in_memory_db) + insert_document(doc_path, config=simple_config) # Update the vector index with the new document. - update_vector_index(config=in_memory_db) + update_vector_index(config=simple_config) # Search for a query. query = "What does it mean for two events to be simultaneous?" - chunk_rowids, scores = hybrid_search(query, config=in_memory_db) + chunk_rowids, scores = hybrid_search(query, config=simple_config) assert len(chunk_rowids) == len(scores) assert all(isinstance(rowid, int) for rowid in chunk_rowids) assert all(isinstance(score, float) for score in scores) # Group the chunks into segments and retrieve them. - segments = retrieve_segments(chunk_rowids, neighbors=None, config=in_memory_db) + segments = retrieve_segments(chunk_rowids, neighbors=None, config=simple_config) assert all(isinstance(segment, str) for segment in segments) assert "Definition of Simultaneity" in segments[0]