@@ -30,7 +30,8 @@ def generate_dummy_documents(count: int) -> List[Dict[str, Any]]:
3030"id" :i ,
3131"text" :"This is a test document: {}" .format (i ),
3232"some_random_thing" :"This will be metadata on it" ,
33- "metadata" : {"uuid" :i * 10 ,"name" :"Test Document {}" .format (i )},
33+ "uuid" :i * 10 ,
34+ "name" :"Test Document {}" .format (i ),
3435 }
3536 )
3637return dummy_documents
@@ -135,6 +136,31 @@ async def test_can_vector_search_with_query_builder_with_remote_embeddings():
135136await collection .archive ()
136137
137138
139+ @pytest .mark .asyncio
140+ async def test_can_vector_search_with_query_builder_and_metadata_filtering ():
141+ model = pgml .Model ()
142+ splitter = pgml .Splitter ()
143+ pipeline = pgml .Pipeline ("test_p_p_tcvswqbamf_1" ,model ,splitter )
144+ collection = pgml .Collection (name = "test_p_c_tcvswqbamf_2" )
145+ await collection .upsert_documents (generate_dummy_documents (3 ))
146+ await collection .add_pipeline (pipeline )
147+ results = (
148+ await collection .query ()
149+ .vector_recall ("Here is some query" ,pipeline )
150+ .filter ({
151+ "metadata" : {
152+ "uuid" : {
153+ "$eq" :0
154+ }
155+ }
156+ })
157+ .limit (10 )
158+ .fetch_all ()
159+ )
160+ assert len (results )== 1
161+ await collection .archive ()
162+
163+
138164###################################################
139165## Test user output facing functions ##############
140166###################################################
@@ -158,17 +184,17 @@ async def test_pipeline_to_dict():
158184###################################################
159185
160186
161- def vector_search (collection_name ,pipeline_name ):
162- collection = pgml .Collection (collection_name )
163- pipeline = pgml .Pipeline (pipeline_name )
164- result = asyncio .run (
165- collection .query ()
166- .vector_recall ("Here is some query" ,pipeline )
167- .limit (10 )
168- .fetch_all ()
169- )
170- print (result )
171- return [0 ,1 ,2 ]
187+ # def vector_search(collection_name, pipeline_name):
188+ # collection = pgml.Collection(collection_name)
189+ # pipeline = pgml.Pipeline(pipeline_name)
190+ # result = asyncio.run(
191+ # collection.query()
192+ # .vector_recall("Here is some query", pipeline)
193+ # .limit(10)
194+ # .fetch_all()
195+ # )
196+ # print(result)
197+ # return [0, 1, 2]
172198
173199
174200# @pytest.mark.asyncio
@@ -200,23 +226,23 @@ def vector_search(collection_name, pipeline_name):
200226###################################################
201227
202228
203- async def silas_test_add_pipeline ():
204- model = pgml .Model ()
205- splitter = pgml .Splitter ()
206- pipeline = pgml .Pipeline ("silas_test_p_1" ,model ,splitter )
207- collection = pgml .Collection (name = "silas_test_c_10" )
208- await collection .add_pipeline (pipeline )
209-
210- async def silas_test_upsert_documents ():
211- collection = pgml .Collection (name = "silas_test_c_9" )
212- await collection .upsert_documents (generate_dummy_documents (10 ))
213-
214- async def silas_test_vector_search ():
215- pipeline = pgml .Pipeline ("silas_test_p_1" )
216- collection = pgml .Collection (name = "silas_test_c_9" )
217- results = await collection .vector_search ("Here is some query" ,pipeline )
218- print (results )
219-
220- # asyncio.run(silas_test_add_pipeline ())
221- # asyncio.run(silas_test_upsert_documents ())
222- # asyncio.run(silas_test_vector_search ())
229+ # async deftest_add_pipeline ():
230+ # model = pgml.Model()
231+ # splitter = pgml.Splitter()
232+ # pipeline = pgml.Pipeline("silas_test_p_1", model, splitter)
233+ # collection = pgml.Collection(name="silas_test_c_10")
234+ # await collection.add_pipeline(pipeline)
235+ #
236+ # async deftest_upsert_documents ():
237+ # collection = pgml.Collection(name="silas_test_c_9")
238+ # await collection.upsert_documents(generate_dummy_documents(10))
239+ #
240+ # async deftest_vector_search ():
241+ # pipeline = pgml.Pipeline("silas_test_p_1")
242+ # collection = pgml.Collection(name="silas_test_c_9")
243+ # results = await collection.vector_search("Here is some query", pipeline)
244+ # print(results)
245+
246+ # asyncio.run(test_add_pipeline ())
247+ # asyncio.run(test_upsert_documents ())
248+ # asyncio.run(test_vector_search ())