@@ -91,7 +91,7 @@ class MockEmbeddingClient(BaseEmbeddingClient):
9191 def __init__ (self ) -> None :
9292 super ().__init__ ()
9393 self .values : list [Any ] = []
94- self .options : EmbeddingGenerationOptions | None = None
94+ self .options : dict [ str , Any ] | None = None
9595
9696 async def get_embeddings (
9797 self ,
@@ -100,7 +100,7 @@ async def get_embeddings(
100100 options : EmbeddingGenerationOptions | None = None ,
101101 ) -> GeneratedEmbeddings [list [float ]]:
102102 self .values = list (values )
103- self .options = options
103+ self .options = dict ( options ) if options is not None else None
104104 return GeneratedEmbeddings ([Embedding (vector = [float (len (str (value ))), 0.5 ]) for value in values ])
105105
106106
@@ -737,6 +737,83 @@ async def test_collection_serializes_records_and_generates_vectors() -> None:
737737 assert embedding_client .options == {"dimensions" : 2 }
738738
739739
740+ async def test_upsert_embedding_options_are_merged_with_field_dimensions () -> None :
741+ embedding_client = MockEmbeddingClient ()
742+ collection = MockCollection (embedding_generator = embedding_client )
743+ caller_options : dict [str , Any ] = {
744+ "task_type" : "RETRIEVAL_DOCUMENT" ,
745+ "extra_parameters" : {"provider_flag" : True },
746+ "dimensions" : 2 ,
747+ }
748+
749+ await collection .upsert ([Record ("one" , "first" , "document" )], embeddings_options = caller_options )
750+
751+ assert embedding_client .options == caller_options
752+ assert embedding_client .values == ["document" ]
753+ assert embedding_client .options is not None
754+ embedding_client .options ["extra_parameters" ]["provider_flag" ] = False
755+ assert caller_options ["extra_parameters" ] == {"provider_flag" : True }
756+ assert collection .records ["one" ]["vector" ] == [8.0 , 0.5 ]
757+
758+ embedding_client .values .clear ()
759+ with pytest .raises (ValueError , match = "must match its declared dimensions" ):
760+ await collection .upsert ([Record ("two" , "second" , "document" )], embeddings_options = {"dimensions" : 3 })
761+ assert embedding_client .values == []
762+ assert "two" not in collection .records
763+
764+
765+ @pytest .mark .parametrize ("invalid_dimensions" , [True , 2.0 , None ])
766+ async def test_upsert_rejects_invalid_dimensions_even_if_equal (invalid_dimensions : Any ) -> None :
767+ embedding_client = MockEmbeddingClient ()
768+ collection = MockCollection (embedding_generator = embedding_client )
769+ with pytest .raises (ValueError , match = "must match its declared dimensions" ):
770+ await collection .upsert (
771+ [Record ("one" , "body" , "document" )], embeddings_options = {"dimensions" : invalid_dimensions }
772+ )
773+ assert embedding_client .values == []
774+
775+
776+ async def test_upsert_field_options_are_isolated_and_prevalidated () -> None :
777+ first_client = MockEmbeddingClient ()
778+ second_client = MockEmbeddingClient ()
779+ definition = VectorStoreCollectionDefinition ([
780+ VectorStoreField ("key" , name = "id" ),
781+ VectorStoreField ("vector" , name = "first" , dimensions = 2 , embedding_generator = first_client ),
782+ VectorStoreField ("vector" , name = "second" , dimensions = 2 , embedding_generator = second_client ),
783+ ])
784+ handler = VectorStoreRecordHandler (dict , definition = definition )
785+ record = {"id" : "one" , "first" : "first text" , "second" : "second text" }
786+
787+ await handler .serialize (record , embeddings_options = {"task_type" : "RETRIEVAL_DOCUMENT" })
788+ assert first_client .options == {"task_type" : "RETRIEVAL_DOCUMENT" , "dimensions" : 2 }
789+ assert second_client .options == {"task_type" : "RETRIEVAL_DOCUMENT" , "dimensions" : 2 }
790+
791+ by_field = {"first" : {"task_type" : "RETRIEVAL_DOCUMENT" }, "second" : {"task_type" : "RETRIEVAL_QUERY" }}
792+ await handler .serialize (record , embeddings_options_by_field = by_field )
793+ assert first_client .options == {"task_type" : "RETRIEVAL_DOCUMENT" , "dimensions" : 2 }
794+ assert second_client .options == {"task_type" : "RETRIEVAL_QUERY" , "dimensions" : 2 }
795+ assert by_field == {"first" : {"task_type" : "RETRIEVAL_DOCUMENT" }, "second" : {"task_type" : "RETRIEVAL_QUERY" }}
796+
797+ first_client .values .clear ()
798+ second_client .values .clear ()
799+ with pytest .raises (ValueError , match = "second.*declared dimensions" ):
800+ await handler .serialize (record , embeddings_options_by_field = {"second" : {"dimensions" : 3 }})
801+ assert first_client .values == second_client .values == []
802+
803+ with pytest .raises (ValueError , match = "either embeddings_options or embeddings_options_by_field" ):
804+ await handler .serialize (record , embeddings_options = {}, embeddings_options_by_field = {})
805+ with pytest .raises (ValueError , match = "non-generated vector field" ):
806+ await handler .serialize (record , embeddings_options_by_field = {"missing" : {}})
807+ with pytest .raises (ValueError , match = "non-generated vector field" ):
808+ await handler .serialize (
809+ record ,
810+ generate_vectors = ["first" ],
811+ embeddings_options_by_field = {"second" : {}},
812+ )
813+ with pytest .raises (ValueError , match = "require records and generated vector fields" ):
814+ await handler .serialize (record , generate_vectors = False , embeddings_options = {})
815+
816+
740817async def test_collection_empty_upsert_skips_embedding_generation () -> None :
741818 embedding_client = MockEmbeddingClient ()
742819 collection = MockCollection (embedding_generator = embedding_client )
@@ -1053,6 +1130,36 @@ async def test_vector_search_generates_query_vector_and_forwards_threshold() ->
10531130 assert responses [1 ]["score" ] == 0.4
10541131
10551132
1133+ async def test_vector_search_uses_provider_options_and_declared_dimensions () -> None :
1134+ embedding_client = MockEmbeddingClient ()
1135+ collection = MockCollection (embedding_generator = embedding_client )
1136+ options = {"task_type" : "RETRIEVAL_QUERY" }
1137+
1138+ await collection .search ("find this" , embeddings_options = options )
1139+
1140+ assert embedding_client .options == {"task_type" : "RETRIEVAL_QUERY" , "dimensions" : 2 }
1141+ assert options == {"task_type" : "RETRIEVAL_QUERY" }
1142+ assert collection .last_search_vector == [9.0 , 0.5 ]
1143+
1144+ embedding_client .values .clear ()
1145+ with pytest .raises (ValueError , match = "must match its declared dimensions" ):
1146+ await collection .search ("find this" , embeddings_options = {"dimensions" : 3 })
1147+ assert embedding_client .values == []
1148+
1149+ await collection .search (vector = [1.0 , 0.0 ], embeddings_options = {"dimensions" : 3 , ** options })
1150+ assert collection .last_search_vector == [1.0 , 0.0 ]
1151+ assert embedding_client .values == []
1152+
1153+ await collection .search ("keyword terms" , vector = [0.0 , 1.0 ], embeddings_options = {"dimensions" : 3 })
1154+ assert collection .last_search_values == "keyword terms"
1155+ assert collection .last_search_vector == [0.0 , 1.0 ]
1156+ assert embedding_client .values == []
1157+
1158+ await MockCollection ().search (vector = [1.0 , 0.0 ], embeddings_options = options )
1159+ with pytest .raises (ValueError , match = "requires an embedding generator" ):
1160+ await MockCollection ().search ("find this" , embeddings_options = options )
1161+
1162+
10561163async def test_keyword_hybrid_search_uses_single_search_method () -> None :
10571164 collection = MockCollection ()
10581165
@@ -1333,6 +1440,44 @@ async def test_create_search_tool_returns_mapped_results() -> None:
13331440 assert result [0 ].text == "one:0.9"
13341441
13351442
1443+ @pytest .mark .parametrize ("search_type" , ["vector" , "keyword_hybrid" ])
1444+ async def test_create_search_tool_forwards_embedding_options (search_type : SearchType ) -> None :
1445+ embedding_client = MockEmbeddingClient ()
1446+ collection = MockCollection (embedding_generator = embedding_client )
1447+ options = {"task_type" : "RETRIEVAL_QUERY" }
1448+
1449+ tool = create_vector_search_tool (
1450+ collection ,
1451+ search_type = search_type ,
1452+ embeddings_options = options ,
1453+ top = 1 ,
1454+ )
1455+ options ["task_type" ] = "RETRIEVAL_DOCUMENT"
1456+
1457+ await tool (query = "find this" )
1458+
1459+ assert embedding_client .values == ["find this" ]
1460+ assert embedding_client .options == {"task_type" : "RETRIEVAL_QUERY" , "dimensions" : 2 }
1461+ assert collection .last_search_values == "find this"
1462+ assert collection .last_search_vector == [9.0 , 0.5 ]
1463+ assert collection .last_search_type == search_type
1464+ assert set (tool .parameters ()["properties" ]) == {"query" }
1465+
1466+
1467+ async def test_create_search_tool_rejects_unusable_embedding_options () -> None :
1468+ tool = create_vector_search_tool (MockCollection (), embeddings_options = {"task_type" : "RETRIEVAL_QUERY" })
1469+ with pytest .raises (ValueError , match = "requires an embedding generator" ):
1470+ await tool (query = "find this" )
1471+
1472+ embedding_client = MockEmbeddingClient ()
1473+ collection = MockCollection (embedding_generator = embedding_client )
1474+ tool = create_vector_search_tool (collection , embeddings_options = {"dimensions" : 3 })
1475+ with pytest .raises (ValueError , match = "must match its declared dimensions" ):
1476+ await tool (query = "find this" )
1477+ assert embedding_client .values == []
1478+ assert collection .last_search_type is None
1479+
1480+
13361481async def test_create_search_tool_supports_declared_filter_parameters () -> None :
13371482 collection = MockCollection ()
13381483 collection .records ["one" ] = {"record_id" : "one" , "body" : "first" , "vector" : [1.0 , 0.0 ]}
@@ -2595,6 +2740,33 @@ async def test_vector_crud_tools_round_trip_records() -> None:
25952740 assert await get_tool .invoke (arguments = {"keys" : ["one" ]}, skip_parsing = True ) == {"records" : []}
25962741
25972742
2743+ async def test_upsert_tool_forwards_provider_embedding_options () -> None :
2744+ embedding_client = MockEmbeddingClient ()
2745+ collection = MockCollection (embedding_generator = embedding_client )
2746+ options = {"task_type" : "RETRIEVAL_DOCUMENT" }
2747+ upsert_tool = create_upsert_tool (collection , embeddings_options = options )
2748+ options ["task_type" ] = "RETRIEVAL_QUERY"
2749+
2750+ await upsert_tool .invoke (
2751+ arguments = {"records" : [{"id" : "one" , "text" : "body" , "vector" : "document" }]},
2752+ skip_parsing = True ,
2753+ )
2754+
2755+ assert embedding_client .options == {"task_type" : "RETRIEVAL_DOCUMENT" , "dimensions" : 2 }
2756+ assert collection .records ["one" ]["vector" ] == [8.0 , 0.5 ]
2757+ with pytest .raises (ValueError , match = "either embeddings_options or embeddings_options_by_field" ):
2758+ create_upsert_tool (collection , embeddings_options = {}, embeddings_options_by_field = {})
2759+
2760+ by_field = {"vector" : {"task_type" : "RETRIEVAL_DOCUMENT" }}
2761+ by_field_tool = create_upsert_tool (collection , embeddings_options_by_field = by_field )
2762+ by_field ["vector" ]["task_type" ] = "RETRIEVAL_QUERY"
2763+ await by_field_tool .invoke (
2764+ arguments = {"records" : [{"id" : "two" , "text" : "body" , "vector" : "more text" }]},
2765+ skip_parsing = True ,
2766+ )
2767+ assert embedding_client .options == {"task_type" : "RETRIEVAL_DOCUMENT" , "dimensions" : 2 }
2768+
2769+
25982770async def test_vector_crud_tools_support_auto_generated_keys_when_model_can_omit_them () -> None :
25992771 definition = VectorStoreCollectionDefinition (
26002772 [
@@ -2740,6 +2912,36 @@ def test_vector_collection_context_provider_configures_tools_and_approvals() ->
27402912 assert all (tool .approval_mode == "always_require" for tool in require_all .tools )
27412913
27422914
2915+ async def test_vector_collection_context_provider_forwards_embedding_options () -> None :
2916+ embedding_client = MockEmbeddingClient ()
2917+ collection = MockCollection (embedding_generator = embedding_client )
2918+ provider = VectorCollectionContextProvider (
2919+ collection ,
2920+ scope_filter = None ,
2921+ upsert_embeddings_options = {"task_type" : "RETRIEVAL_DOCUMENT" },
2922+ search_embeddings_options = {"task_type" : "RETRIEVAL_QUERY" },
2923+ )
2924+ tools = {tool .name : tool for tool in provider .tools }
2925+
2926+ await tools ["upsert" ].invoke (
2927+ arguments = {"records" : [{"id" : "one" , "text" : "body" , "vector" : "document" }]},
2928+ skip_parsing = True ,
2929+ )
2930+ assert embedding_client .options == {"task_type" : "RETRIEVAL_DOCUMENT" , "dimensions" : 2 }
2931+
2932+ await tools ["search" ].invoke (arguments = {"query" : "find this" }, skip_parsing = True )
2933+ assert embedding_client .options == {"task_type" : "RETRIEVAL_QUERY" , "dimensions" : 2 }
2934+
2935+ with pytest .raises (ValueError , match = "Upsert embedding options require include_upsert_tool=True" ):
2936+ VectorCollectionContextProvider (
2937+ collection , scope_filter = None , include_upsert_tool = False , upsert_embeddings_options = {}
2938+ )
2939+ with pytest .raises (ValueError , match = "Search embedding options require include_search_tool=True" ):
2940+ VectorCollectionContextProvider (
2941+ collection , scope_filter = None , include_search_tool = False , search_embeddings_options = {}
2942+ )
2943+
2944+
27432945async def test_vector_collection_context_provider_adds_attributed_context () -> None :
27442946 collection = MockCollection (embedding_generator = MockEmbeddingClient ())
27452947 details_tool = create_vector_search_tool (collection , name = "search_details" )
0 commit comments