mirror of
https://github.com/Mintplex-Labs/langchain-python.git
synced 2026-07-19 21:33:31 -04:00
Update SinglStoreDB vectorstore (#6423)
1. Introduced new distance strategies support: **DOT_PRODUCT** and **EUCLIDEAN_DISTANCE** for enhanced flexibility. 2. Implemented a feature to filter results based on metadata fields. 3. Incorporated connection attributes specifying "langchain python sdk" usage for enhanced traceability and debugging. 4. Expanded the suite of integration tests for improved code reliability. 5. Updated the existing notebook with the usage example @dev2049 --------- Co-authored-by: Volodymyr Tkachuk <vtkachuk-ua@singlestore.com> Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
This commit is contained in:
@@ -5,7 +5,7 @@ import numpy as np
|
||||
import pytest
|
||||
|
||||
from langchain.docstore.document import Document
|
||||
from langchain.vectorstores.singlestoredb import SingleStoreDB
|
||||
from langchain.vectorstores.singlestoredb import DistanceStrategy, SingleStoreDB
|
||||
from tests.integration_tests.vectorstores.fake_embeddings import FakeEmbeddings
|
||||
|
||||
TEST_SINGLESTOREDB_URL = "root:pass@localhost:3306/db"
|
||||
@@ -80,6 +80,24 @@ def test_singlestoredb_new_vector(texts: List[str]) -> None:
|
||||
drop(table_name)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not singlestoredb_installed, reason="singlestoredb not installed")
|
||||
def test_singlestoredb_euclidean_distance(texts: List[str]) -> None:
|
||||
"""Test adding a new document"""
|
||||
table_name = "test_singlestoredb_euclidean_distance"
|
||||
drop(table_name)
|
||||
docsearch = SingleStoreDB.from_texts(
|
||||
texts,
|
||||
FakeEmbeddings(),
|
||||
distance_strategy=DistanceStrategy.EUCLIDEAN_DISTANCE,
|
||||
table_name=table_name,
|
||||
host=TEST_SINGLESTOREDB_URL,
|
||||
)
|
||||
docsearch.add_texts(["foo"])
|
||||
output = docsearch.similarity_search("foo", k=2)
|
||||
assert output == TEST_RESULT
|
||||
drop(table_name)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not singlestoredb_installed, reason="singlestoredb not installed")
|
||||
def test_singlestoredb_from_existing(texts: List[str]) -> None:
|
||||
"""Test adding a new document"""
|
||||
@@ -140,3 +158,193 @@ def test_singlestoredb_add_texts_to_existing(texts: List[str]) -> None:
|
||||
output = docsearch.similarity_search("foo", k=2)
|
||||
assert output == TEST_RESULT
|
||||
drop(table_name)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not singlestoredb_installed, reason="singlestoredb not installed")
|
||||
def test_singlestoredb_filter_metadata(texts: List[str]) -> None:
|
||||
"""Test filtering by metadata"""
|
||||
table_name = "test_singlestoredb_filter_metadata"
|
||||
drop(table_name)
|
||||
docs = [
|
||||
Document(page_content=t, metadata={"index": i}) for i, t in enumerate(texts)
|
||||
]
|
||||
docsearch = SingleStoreDB.from_documents(
|
||||
docs,
|
||||
FakeEmbeddings(),
|
||||
distance_strategy=DistanceStrategy.EUCLIDEAN_DISTANCE,
|
||||
table_name=table_name,
|
||||
host=TEST_SINGLESTOREDB_URL,
|
||||
)
|
||||
output = docsearch.similarity_search("foo", k=1, filter={"index": 2})
|
||||
assert output == [Document(page_content="baz", metadata={"index": 2})]
|
||||
drop(table_name)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not singlestoredb_installed, reason="singlestoredb not installed")
|
||||
def test_singlestoredb_filter_metadata_2(texts: List[str]) -> None:
|
||||
"""Test filtering by metadata field that is similar for each document"""
|
||||
table_name = "test_singlestoredb_filter_metadata_2"
|
||||
drop(table_name)
|
||||
docs = [
|
||||
Document(page_content=t, metadata={"index": i, "category": "budget"})
|
||||
for i, t in enumerate(texts)
|
||||
]
|
||||
docsearch = SingleStoreDB.from_documents(
|
||||
docs,
|
||||
FakeEmbeddings(),
|
||||
distance_strategy=DistanceStrategy.EUCLIDEAN_DISTANCE,
|
||||
table_name=table_name,
|
||||
host=TEST_SINGLESTOREDB_URL,
|
||||
)
|
||||
output = docsearch.similarity_search("foo", k=1, filter={"category": "budget"})
|
||||
assert output == [
|
||||
Document(page_content="foo", metadata={"index": 0, "category": "budget"})
|
||||
]
|
||||
drop(table_name)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not singlestoredb_installed, reason="singlestoredb not installed")
|
||||
def test_singlestoredb_filter_metadata_3(texts: List[str]) -> None:
|
||||
"""Test filtering by two metadata fields"""
|
||||
table_name = "test_singlestoredb_filter_metadata_3"
|
||||
drop(table_name)
|
||||
docs = [
|
||||
Document(page_content=t, metadata={"index": i, "category": "budget"})
|
||||
for i, t in enumerate(texts)
|
||||
]
|
||||
docsearch = SingleStoreDB.from_documents(
|
||||
docs,
|
||||
FakeEmbeddings(),
|
||||
distance_strategy=DistanceStrategy.EUCLIDEAN_DISTANCE,
|
||||
table_name=table_name,
|
||||
host=TEST_SINGLESTOREDB_URL,
|
||||
)
|
||||
output = docsearch.similarity_search(
|
||||
"foo", k=1, filter={"category": "budget", "index": 1}
|
||||
)
|
||||
assert output == [
|
||||
Document(page_content="bar", metadata={"index": 1, "category": "budget"})
|
||||
]
|
||||
drop(table_name)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not singlestoredb_installed, reason="singlestoredb not installed")
|
||||
def test_singlestoredb_filter_metadata_4(texts: List[str]) -> None:
|
||||
"""Test no matches"""
|
||||
table_name = "test_singlestoredb_filter_metadata_4"
|
||||
drop(table_name)
|
||||
docs = [
|
||||
Document(page_content=t, metadata={"index": i, "category": "budget"})
|
||||
for i, t in enumerate(texts)
|
||||
]
|
||||
docsearch = SingleStoreDB.from_documents(
|
||||
docs,
|
||||
FakeEmbeddings(),
|
||||
distance_strategy=DistanceStrategy.EUCLIDEAN_DISTANCE,
|
||||
table_name=table_name,
|
||||
host=TEST_SINGLESTOREDB_URL,
|
||||
)
|
||||
output = docsearch.similarity_search("foo", k=1, filter={"category": "vacation"})
|
||||
assert output == []
|
||||
drop(table_name)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not singlestoredb_installed, reason="singlestoredb not installed")
|
||||
def test_singlestoredb_filter_metadata_5(texts: List[str]) -> None:
|
||||
"""Test complex metadata path"""
|
||||
table_name = "test_singlestoredb_filter_metadata_5"
|
||||
drop(table_name)
|
||||
docs = [
|
||||
Document(
|
||||
page_content=t,
|
||||
metadata={
|
||||
"index": i,
|
||||
"category": "budget",
|
||||
"subfield": {"subfield": {"idx": i, "other_idx": i + 1}},
|
||||
},
|
||||
)
|
||||
for i, t in enumerate(texts)
|
||||
]
|
||||
docsearch = SingleStoreDB.from_documents(
|
||||
docs,
|
||||
FakeEmbeddings(),
|
||||
distance_strategy=DistanceStrategy.EUCLIDEAN_DISTANCE,
|
||||
table_name=table_name,
|
||||
host=TEST_SINGLESTOREDB_URL,
|
||||
)
|
||||
output = docsearch.similarity_search(
|
||||
"foo", k=1, filter={"category": "budget", "subfield": {"subfield": {"idx": 2}}}
|
||||
)
|
||||
assert output == [
|
||||
Document(
|
||||
page_content="baz",
|
||||
metadata={
|
||||
"index": 2,
|
||||
"category": "budget",
|
||||
"subfield": {"subfield": {"idx": 2, "other_idx": 3}},
|
||||
},
|
||||
)
|
||||
]
|
||||
drop(table_name)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not singlestoredb_installed, reason="singlestoredb not installed")
|
||||
def test_singlestoredb_filter_metadata_6(texts: List[str]) -> None:
|
||||
"""Test filtering by other bool"""
|
||||
table_name = "test_singlestoredb_filter_metadata_6"
|
||||
drop(table_name)
|
||||
docs = [
|
||||
Document(
|
||||
page_content=t,
|
||||
metadata={"index": i, "category": "budget", "is_good": i == 1},
|
||||
)
|
||||
for i, t in enumerate(texts)
|
||||
]
|
||||
docsearch = SingleStoreDB.from_documents(
|
||||
docs,
|
||||
FakeEmbeddings(),
|
||||
distance_strategy=DistanceStrategy.EUCLIDEAN_DISTANCE,
|
||||
table_name=table_name,
|
||||
host=TEST_SINGLESTOREDB_URL,
|
||||
)
|
||||
output = docsearch.similarity_search(
|
||||
"foo", k=1, filter={"category": "budget", "is_good": True}
|
||||
)
|
||||
assert output == [
|
||||
Document(
|
||||
page_content="bar",
|
||||
metadata={"index": 1, "category": "budget", "is_good": True},
|
||||
)
|
||||
]
|
||||
drop(table_name)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not singlestoredb_installed, reason="singlestoredb not installed")
|
||||
def test_singlestoredb_filter_metadata_7(texts: List[str]) -> None:
|
||||
"""Test filtering by float"""
|
||||
table_name = "test_singlestoredb_filter_metadata_7"
|
||||
drop(table_name)
|
||||
docs = [
|
||||
Document(
|
||||
page_content=t,
|
||||
metadata={"index": i, "category": "budget", "score": i + 0.5},
|
||||
)
|
||||
for i, t in enumerate(texts)
|
||||
]
|
||||
docsearch = SingleStoreDB.from_documents(
|
||||
docs,
|
||||
FakeEmbeddings(),
|
||||
distance_strategy=DistanceStrategy.EUCLIDEAN_DISTANCE,
|
||||
table_name=table_name,
|
||||
host=TEST_SINGLESTOREDB_URL,
|
||||
)
|
||||
output = docsearch.similarity_search(
|
||||
"bar", k=1, filter={"category": "budget", "score": 2.5}
|
||||
)
|
||||
assert output == [
|
||||
Document(
|
||||
page_content="baz",
|
||||
metadata={"index": 2, "category": "budget", "score": 2.5},
|
||||
)
|
||||
]
|
||||
drop(table_name)
|
||||
|
||||
Reference in New Issue
Block a user