wehub-resource-sync
121 lines
4.6 KiB
Python
121 lines
4.6 KiB
Python
|
|
""" Provides a set of short tests for specific vector DB. """
|
|
|
|
import os
|
|
import pytest
|
|
import time
|
|
|
|
from llmware.library import Library
|
|
from llmware.embeddings import EmbeddingHandler
|
|
from llmware.exceptions import UnsupportedEmbeddingDatabaseException
|
|
from llmware.models import ModelCatalog
|
|
from llmware.retrieval import Query
|
|
from llmware.setup import Setup
|
|
from llmware.configs import LLMWareConfig
|
|
from llmware.resources import CloudBucketManager
|
|
from tests.embeddings.utils import qdrant_installed
|
|
|
|
|
|
def test_unsupported_embedding_db():
|
|
|
|
embedding_db = "milvusXYZ" # Bad Embedding DB Name
|
|
with pytest.raises(UnsupportedEmbeddingDatabaseException) as excinfo:
|
|
embedding_handler = EmbeddingHandler(library=None)
|
|
embedding_summary = embedding_handler.create_new_embedding(embedding_db=embedding_db, model=None)
|
|
assert str(excinfo.value) == f"'{embedding_db}' is not a supported vector embedding database"
|
|
|
|
|
|
def test_milvus_embedding_and_query():
|
|
|
|
sample_files_path = Setup().load_sample_files()
|
|
library = Library().create_new_library("test_embedding_milvus")
|
|
library.add_files(os.path.join(sample_files_path,"SmallLibrary"))
|
|
results = generic_embedding_and_query(library, "milvus")
|
|
assert len(results) > 0
|
|
library.delete_library(confirm_delete=True)
|
|
|
|
|
|
def test_neo4j_embedding_and_query():
|
|
|
|
sample_files_path = Setup().load_sample_files()
|
|
library = Library().create_new_library("test_embedding_neo4j")
|
|
library.add_files(os.path.join(sample_files_path,"SmallLibrary"))
|
|
results = generic_embedding_and_query(library, "neo4j")
|
|
assert len(results) > 0
|
|
library.delete_library(confirm_delete=True)
|
|
|
|
|
|
def test_chromadb_embedding_and_query():
|
|
|
|
sample_files_path = Setup().load_sample_files()
|
|
library = Library().create_new_library("test_embedding_neo4j")
|
|
library.add_files(os.path.join(sample_files_path,"SmallLibrary"))
|
|
results = generic_embedding_and_query(library, "chromadb")
|
|
assert len(results) > 0
|
|
library.delete_library(confirm_delete=True)
|
|
|
|
|
|
def test_faiss_embedding_and_query():
|
|
|
|
sample_files_path = Setup().load_sample_files()
|
|
library = Library().create_new_library("test_embedding_faiss")
|
|
library.add_files(os.path.join(sample_files_path,"SmallLibrary"))
|
|
results = generic_embedding_and_query(library, "faiss")
|
|
assert len(results) > 0
|
|
library.delete_library(confirm_delete=True)
|
|
|
|
|
|
def test_lancedb_embedding_and_query():
|
|
|
|
sample_files_path = Setup().load_sample_files()
|
|
library = Library().create_new_library("test_embedding_lancedb")
|
|
library.add_files(os.path.join(sample_files_path,"SmallLibrary"))
|
|
results = generic_embedding_and_query(library, "lancedb")
|
|
assert len(results) > 0
|
|
library.delete_library(confirm_delete=True)
|
|
|
|
|
|
@pytest.mark.skipif(not qdrant_installed(), reason="Qdrant client is not installed")
|
|
def test_qdrant_embedding_and_query():
|
|
os.environ["USER_MANAGED_QDRANT_LOCATION"] = ":memory:"
|
|
sample_files_path = Setup().load_sample_files()
|
|
library = Library().create_new_library("test_embedding_qdrant")
|
|
library.add_files(os.path.join(sample_files_path,"SmallLibrary"))
|
|
results = generic_embedding_and_query(library, "qdrant")
|
|
assert len(results) > 0
|
|
library.delete_library(confirm_delete=True)
|
|
|
|
# def test_pinecone_embedding_and_query():
|
|
# with pytest.raises(ImportError) as excinfo:
|
|
# library = None
|
|
# results = generic_embedding_and_query(library, "pinecone")
|
|
# assert 'pip install pinecone-client' in str(excinfo.value)
|
|
|
|
# def test_pinecone_embedding_and_query():
|
|
# sample_files_path = Setup().load_sample_files()
|
|
# library = Library().create_new_library("test_embedding")
|
|
# library.add_files(os.path.join(sample_files_path,"SmallLibrary"))
|
|
|
|
# os.environ["PINECONE_API_KEY"] = "abeab29e-7a48-426b-b17a-c74567720876"
|
|
# os.environ["PINECONE_ENVIRONMENT"] = "gcp-starter"
|
|
|
|
# results = generic_embedding_and_query(library, "pinecone")
|
|
# assert len(results) > 0
|
|
# library.delete_library(confirm_delete=True)
|
|
|
|
|
|
def generic_embedding_and_query(library, embedding_db):
|
|
|
|
# Run the embeddings (only of first 3 docs )
|
|
model=ModelCatalog().load_model("mini-lm-sbert")
|
|
embedding_handler = EmbeddingHandler(library=library)
|
|
embedding_summary = embedding_handler.create_new_embedding(embedding_db=embedding_db, model=model,doc_ids=[1, 2, 3, 4, 5])
|
|
|
|
# Do a query
|
|
query = "pact"
|
|
query_results = Query(library).semantic_query(query, result_count=5)
|
|
|
|
# Delete the embedding
|
|
embedding_handler.delete_index(embedding_db=embedding_db,embedding_dims=embedding_summary["embedding_dims"], model_name=model.model_name)
|
|
return query_results
|