diff --git a/test/collection/test_classes_generative.py b/test/collection/test_classes_generative.py index 5d9496417..cd4a6768d 100644 --- a/test/collection/test_classes_generative.py +++ b/test/collection/test_classes_generative.py @@ -98,14 +98,13 @@ def test_generative_parameters_images_parsing( ), ), ( - GenerativeConfig.aws( + GenerativeConfig.aws_sagemaker( endpoint="http://localhost:8080", - model="text-to-image", region="us-west-2", - service="sagemaker", target_model="arn:aws:sagemaker:us-west-2:123456789012:model/text-to-image", target_variant="variant-1", temperature=0.5, + stop_sequences=["\n"], )._to_grpc( _GenerativeConfigRuntimeOptions( return_metadata=True, images=[LOGO_ENCODED], image_properties=["image"] @@ -115,12 +114,12 @@ def test_generative_parameters_images_parsing( return_metadata=True, aws=generative_pb2.GenerativeAWS( endpoint="http://localhost:8080", - model="text-to-image", region="us-west-2", service="sagemaker", target_model="arn:aws:sagemaker:us-west-2:123456789012:model/text-to-image", target_variant="variant-1", temperature=0.5, + stop_sequences=base_pb2.TextArray(values=["\n"]), images=base_pb2.TextArray(values=[LOGO_ENCODED]), image_properties=base_pb2.TextArray(values=["image"]), ), diff --git a/weaviate/collections/classes/generative.py b/weaviate/collections/classes/generative.py index d5c1a5de2..2a6ec96ee 100644 --- a/weaviate/collections/classes/generative.py +++ b/weaviate/collections/classes/generative.py @@ -130,9 +130,10 @@ def _to_grpc(self, opts: _GenerativeConfigRuntimeOptions) -> generative_pb2.Gene target_variant=self.target_variant, temperature=self.temperature, max_tokens=self.max_tokens, + stop_sequences=_to_text_array(self.stop_sequences), images=_to_text_array(opts.images), image_properties=_to_text_array(opts.image_properties), - # TODO - add top_k, top_p & stop_sequences here when added to server-side proto + # TODO - add top_k & top_p here when added to server-side proto # Check the latest availble version of `grpc/proto/v1/generative.proto` (see GenerativeAWS) in the server repo ), )