Skip to content

Commit 702d6dd

Browse files
committed
test: add hybrid search index related tests
- add tests for creating and validating hybrid search indexes using HybridSearchConfig - add tests for applying hybrid search indexes to tables with and without TSV columns - add tests for applying hybrid search indexes in both async and sync versions - add hybrid table cleanup logic to existing tests
1 parent f532396 commit 702d6dd

File tree

1 file changed

+116
-3
lines changed

1 file changed

+116
-3
lines changed

tests/unit_tests/v2/test_pg_vectorstore_index.py

Lines changed: 116 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from sqlalchemy import text
1010

1111
from langchain_postgres import PGEngine, PGVectorStore
12+
from langchain_postgres.v2.hybrid_search_config import HybridSearchConfig
1213
from langchain_postgres.v2.indexes import (
1314
DistanceStrategy,
1415
HNSWIndex,
@@ -17,12 +18,14 @@
1718
from tests.utils import VECTORSTORE_CONNECTION_STRING as CONNECTION_STRING
1819

1920
uuid_str = str(uuid.uuid4()).replace("-", "_")
20-
uuid_str_sync = str(uuid.uuid4()).replace("-", "_")
21+
uuid_str_async = str(uuid.uuid4()).replace("-", "_")
2122
DEFAULT_TABLE = "default" + uuid_str
22-
DEFAULT_TABLE_ASYNC = "default_sync" + uuid_str_sync
23+
DEFAULT_HYBRID_TABLE = "hybrid" + uuid_str
24+
DEFAULT_TABLE_ASYNC = "default_async" + uuid_str_async
25+
DEFAULT_HYBRID_TABLE_ASYNC = "hybrid_async" + uuid_str_async
2326
CUSTOM_TABLE = "custom" + str(uuid.uuid4()).replace("-", "_")
2427
DEFAULT_INDEX_NAME = "index" + uuid_str
25-
DEFAULT_INDEX_NAME_ASYNC = "index" + uuid_str_sync
28+
DEFAULT_INDEX_NAME_ASYNC = "index" + uuid_str_async
2629
VECTOR_SIZE = 768
2730

2831
embeddings_service = DeterministicFakeEmbedding(size=VECTOR_SIZE)
@@ -64,6 +67,7 @@ async def engine(self) -> AsyncIterator[PGEngine]:
6467
engine = PGEngine.from_connection_string(url=CONNECTION_STRING)
6568
yield engine
6669
await aexecute(engine, f"DROP TABLE IF EXISTS {DEFAULT_TABLE}")
70+
await aexecute(engine, f"DROP TABLE IF EXISTS {DEFAULT_HYBRID_TABLE}")
6771
await engine.close()
6872

6973
@pytest_asyncio.fixture(scope="class")
@@ -118,6 +122,60 @@ async def test_is_valid_index(self, vs: PGVectorStore) -> None:
118122
is_valid = vs.is_valid_index("invalid_index")
119123
assert not is_valid
120124

125+
async def test_apply_hybrid_search_index_non_hybrid_search_vs(
126+
self, vs: PGVectorStore
127+
) -> None:
128+
with pytest.raises(ValueError):
129+
vs.apply_hybrid_search_index()
130+
131+
async def test_apply_hybrid_search_index_table_without_tsv_column(
132+
self, engine: PGEngine, vs: PGVectorStore
133+
) -> None:
134+
tsv_index_name = "tsv_index_on_table_without_tsv_column_" + uuid_str
135+
vs_hybrid = PGVectorStore.create_sync(
136+
engine,
137+
embedding_service=embeddings_service,
138+
table_name=DEFAULT_TABLE,
139+
hybrid_search_config=HybridSearchConfig(index_name=tsv_index_name),
140+
)
141+
is_valid_index = vs_hybrid.is_valid_index(tsv_index_name)
142+
assert is_valid_index == False
143+
vs_hybrid.apply_hybrid_search_index()
144+
assert vs_hybrid.is_valid_index(tsv_index_name)
145+
vs_hybrid.drop_vector_index(tsv_index_name)
146+
is_valid_index = vs_hybrid.is_valid_index(tsv_index_name)
147+
assert is_valid_index == False
148+
149+
async def test_apply_hybrid_search_index_table_with_tsv_column(
150+
self, engine: PGEngine
151+
) -> None:
152+
tsv_index_name = "tsv_index_on_table_with_tsv_column_" + uuid_str
153+
config = HybridSearchConfig(
154+
tsv_column="tsv_column",
155+
tsv_lang="pg_catalog.english",
156+
index_name=tsv_index_name,
157+
)
158+
engine.init_vectorstore_table(
159+
DEFAULT_HYBRID_TABLE,
160+
VECTOR_SIZE,
161+
hybrid_search_config=config,
162+
)
163+
vs_hybrid = PGVectorStore.create_sync(
164+
engine,
165+
embedding_service=embeddings_service,
166+
table_name=DEFAULT_HYBRID_TABLE,
167+
hybrid_search_config=config,
168+
)
169+
is_valid_index = vs_hybrid.is_valid_index(tsv_index_name)
170+
assert is_valid_index == False
171+
vs_hybrid.apply_hybrid_search_index()
172+
assert vs_hybrid.is_valid_index(tsv_index_name)
173+
vs_hybrid.reindex(tsv_index_name)
174+
assert vs_hybrid.is_valid_index(tsv_index_name)
175+
vs_hybrid.drop_vector_index(tsv_index_name)
176+
is_valid_index = vs_hybrid.is_valid_index(tsv_index_name)
177+
assert is_valid_index == False
178+
121179

122180
@pytest.mark.enable_socket
123181
@pytest.mark.asyncio(scope="class")
@@ -127,6 +185,7 @@ async def engine(self) -> AsyncIterator[PGEngine]:
127185
engine = PGEngine.from_connection_string(url=CONNECTION_STRING)
128186
yield engine
129187
await aexecute(engine, f"DROP TABLE IF EXISTS {DEFAULT_TABLE_ASYNC}")
188+
await aexecute(engine, f"DROP TABLE IF EXISTS {DEFAULT_HYBRID_TABLE_ASYNC}")
130189
await engine.close()
131190

132191
@pytest_asyncio.fixture(scope="class")
@@ -179,3 +238,57 @@ async def test_aapply_vector_index_ivfflat(self, vs: PGVectorStore) -> None:
179238
async def test_is_valid_index(self, vs: PGVectorStore) -> None:
180239
is_valid = await vs.ais_valid_index("invalid_index")
181240
assert not is_valid
241+
242+
async def test_aapply_hybrid_search_index_non_hybrid_search_vs(
243+
self, vs: PGVectorStore
244+
) -> None:
245+
with pytest.raises(ValueError):
246+
await vs.aapply_hybrid_search_index()
247+
248+
async def test_aapply_hybrid_search_index_table_without_tsv_column(
249+
self, engine: PGEngine, vs: PGVectorStore
250+
) -> None:
251+
tsv_index_name = "tsv_index_on_table_without_tsv_column_" + uuid_str_async
252+
vs_hybrid = await PGVectorStore.create(
253+
engine,
254+
embedding_service=embeddings_service,
255+
table_name=DEFAULT_TABLE_ASYNC,
256+
hybrid_search_config=HybridSearchConfig(index_name=tsv_index_name),
257+
)
258+
is_valid_index = await vs_hybrid.ais_valid_index(tsv_index_name)
259+
assert is_valid_index == False
260+
await vs_hybrid.aapply_hybrid_search_index()
261+
assert await vs_hybrid.ais_valid_index(tsv_index_name)
262+
await vs_hybrid.adrop_vector_index(tsv_index_name)
263+
is_valid_index = await vs_hybrid.ais_valid_index(tsv_index_name)
264+
assert is_valid_index == False
265+
266+
async def test_aapply_hybrid_search_index_table_with_tsv_column(
267+
self, engine: PGEngine
268+
) -> None:
269+
tsv_index_name = "tsv_index_on_table_with_tsv_column_" + uuid_str_async
270+
config = HybridSearchConfig(
271+
tsv_column="tsv_column",
272+
tsv_lang="pg_catalog.english",
273+
index_name=tsv_index_name,
274+
)
275+
await engine.ainit_vectorstore_table(
276+
DEFAULT_HYBRID_TABLE_ASYNC,
277+
VECTOR_SIZE,
278+
hybrid_search_config=config,
279+
)
280+
vs_hybrid = await PGVectorStore.create(
281+
engine,
282+
embedding_service=embeddings_service,
283+
table_name=DEFAULT_HYBRID_TABLE_ASYNC,
284+
hybrid_search_config=config,
285+
)
286+
is_valid_index = await vs_hybrid.ais_valid_index(tsv_index_name)
287+
assert is_valid_index == False
288+
await vs_hybrid.aapply_hybrid_search_index()
289+
assert await vs_hybrid.ais_valid_index(tsv_index_name)
290+
await vs_hybrid.areindex(tsv_index_name)
291+
assert await vs_hybrid.ais_valid_index(tsv_index_name)
292+
await vs_hybrid.adrop_vector_index(tsv_index_name)
293+
is_valid_index = await vs_hybrid.ais_valid_index(tsv_index_name)
294+
assert is_valid_index == False

0 commit comments

Comments
 (0)