Merge pull request #1819 from weaviate/max_tokens_aws

Add max token to AWS generative module

Author: dirkkul

test/collection/test_config.py

@@ -2060,7 +2060,7 @@ def test_config_with_named_vectors(
     (
         [
             Configure.Vectors.text2vec_aws(
-                name="test", region="us-east-1", source_properties=["prop"]
+                name="test", region="us-east-1", source_properties=["prop"], model="model"
             )
         ],
         {
@@ -2071,6 +2071,7 @@ def test_config_with_named_vectors(
                         "vectorizeClassName": True,
                         "region": "us-east-1",
                         "service": "bedrock",
+                        "model": "model",
                     }
                 },
                 "vectorIndexType": "hnsw",

weaviate/collections/classes/config.py

@@ -464,6 +464,7 @@ class _GenerativeAWSConfig(_GenerativeProvider):
     service: str
     model: Optional[str]
     endpoint: Optional[str]
+    maxTokens: Optional[int]
 
 
 class _GenerativeAnthropicConfig(_GenerativeProvider):
@@ -891,6 +892,7 @@ def aws(
         region: str = "",  # cant have a non-default value after a default value, but we cant change the order for BC
         endpoint: Optional[str] = None,
         service: Union[AWSService, str] = "bedrock",
+        max_tokens: Optional[int] = None,
     ) -> _GenerativeProvider:
         """Create a `_GenerativeAWSConfig` object for use when performing AI generation using the `generative-aws` module.
 
@@ -899,15 +901,13 @@ def aws(
 
         Args:
             model: The model to use, REQUIRED for service "bedrock".
+            max_tokens: The maximum number of tokens to generate. Defaults to `None`, which uses the server-defined default.
             region: The AWS region to run the model from, REQUIRED.
             endpoint: The model to use, REQUIRED for service "sagemaker".
             service: The AWS service to use, options are "bedrock" and "sagemaker".
         """
         return _GenerativeAWSConfig(
-            model=model,
-            region=region,
-            service=service,
-            endpoint=endpoint,
+            model=model, region=region, service=service, endpoint=endpoint, maxTokens=max_tokens
         )
 
     @staticmethod

weaviate/collections/classes/config_vectors.py

@@ -652,7 +652,7 @@ def text2vec_aws(
         name: Optional[str] = None,
         quantizer: Optional[_QuantizerConfigCreate] = None,
         endpoint: Optional[str] = None,
-        model: Optional[Union[AWSModel, str]] = None,
+        model: Optional[Union[AWSModel, str]],
         region: str,
         service: Union[AWSService, str] = "bedrock",
         source_properties: Optional[List[str]] = None,
@@ -668,7 +668,7 @@ def text2vec_aws(
             name: The name of the vector.
             quantizer: The quantizer to use for the vector index. If not provided, no quantization will be applied.
             endpoint: The endpoint to use. Defaults to `None`, which uses the server-defined default.
-            model: The model to use.
+            model: The model to use, REQUIRED.
             region: The AWS region to run the model from, REQUIRED.
             service: The AWS service to use. Defaults to `bedrock`.
             source_properties: Which properties should be included when vectorizing. By default all text properties are included.

weaviate/collections/classes/generative.py

@@ -103,6 +103,7 @@ class _GenerativeAWS(_GenerativeConfigRuntime):
     generative: Union[GenerativeSearches, _EnumLikeStr] = Field(
         default=GenerativeSearches.AWS, frozen=True, exclude=True
     )
+    max_tokens: Optional[int]
     model: Optional[str]
     region: Optional[str]
     endpoint: Optional[AnyHttpUrl]
@@ -122,6 +123,7 @@ def _to_grpc(self, opts: _GenerativeConfigRuntimeOptions) -> generative_pb2.Gene
                 target_model=self.target_model,
                 target_variant=self.target_variant,
                 temperature=self.temperature,
+                max_tokens=self.max_tokens,
                 images=_to_text_array(opts.images),
                 image_properties=_to_text_array(opts.image_properties),
             ),
@@ -474,6 +476,7 @@ def anyscale(
     def aws(
         *,
         endpoint: Optional[str] = None,
+        max_tokens: Optional[int] = None,
         model: Optional[str] = None,
         region: Optional[str] = None,
         service: Optional[Union[AWSService, str]] = None,
@@ -488,6 +491,7 @@ def aws(
 
         Args:
             endpoint: The endpoint to use when requesting the generation. Defaults to `None`, which uses the server-defined default
+            max_tokens: The maximum number of tokens to generate. Defaults to `None`, which uses the server-defined default
             model: The model to use. Defaults to `None`, which uses the server-defined default
             region: The AWS region to run the model from. Defaults to `None`, which uses the server-defined default
             service: The AWS service to use. Defaults to `None`, which uses the server-defined default
@@ -497,6 +501,7 @@ def aws(
         """
         return _GenerativeAWS(
             model=model,
+            max_tokens=max_tokens,
             region=region,
             service=service,
             endpoint=AnyUrl(endpoint) if endpoint is not None else None,

weaviate/collections/classes/internal.py

@@ -306,7 +306,7 @@ def __init__(
     def to_grpc(self, server_version: _ServerVersion) -> generative_pb2.GenerativeSearch:
         if server_version.is_lower_than(1, 27, 14):
             if self.generative_provider is not None:
-                raise WeaviateUnsupportedFeatureError("Dynamic RAG", str(server_version), "1.30.0")
+                raise WeaviateUnsupportedFeatureError("Dynamic RAG", str(server_version), "1.27.14")
 
             if isinstance(self.single, _SinglePrompt):
                 single_prompt: Optional[str] = self.single.prompt

weaviate/proto/v1/v4216/v1/generative_pb2.py

@@ -105,7 +105,7 @@ class GenerativeAnyscale(_message.Message):
     def __init__(self, base_url: _Optional[str] = ..., model: _Optional[str] = ..., temperature: _Optional[float] = ...) -> None: ...
 
 class GenerativeAWS(_message.Message):
-    __slots__ = ["model", "temperature", "service", "region", "endpoint", "target_model", "target_variant", "images", "image_properties"]
+    __slots__ = ["model", "temperature", "service", "region", "endpoint", "target_model", "target_variant", "images", "image_properties", "max_tokens"]
     MODEL_FIELD_NUMBER: _ClassVar[int]
     TEMPERATURE_FIELD_NUMBER: _ClassVar[int]
     SERVICE_FIELD_NUMBER: _ClassVar[int]
@@ -115,6 +115,7 @@ class GenerativeAWS(_message.Message):
     TARGET_VARIANT_FIELD_NUMBER: _ClassVar[int]
     IMAGES_FIELD_NUMBER: _ClassVar[int]
     IMAGE_PROPERTIES_FIELD_NUMBER: _ClassVar[int]
+    MAX_TOKENS_FIELD_NUMBER: _ClassVar[int]
     model: str
     temperature: float
     service: str
@@ -124,7 +125,8 @@ class GenerativeAWS(_message.Message):
     target_variant: str
     images: _base_pb2.TextArray
     image_properties: _base_pb2.TextArray
-    def __init__(self, model: _Optional[str] = ..., temperature: _Optional[float] = ..., service: _Optional[str] = ..., region: _Optional[str] = ..., endpoint: _Optional[str] = ..., target_model: _Optional[str] = ..., target_variant: _Optional[str] = ..., images: _Optional[_Union[_base_pb2.TextArray, _Mapping]] = ..., image_properties: _Optional[_Union[_base_pb2.TextArray, _Mapping]] = ...) -> None: ...
+    max_tokens: int
+    def __init__(self, model: _Optional[str] = ..., temperature: _Optional[float] = ..., service: _Optional[str] = ..., region: _Optional[str] = ..., endpoint: _Optional[str] = ..., target_model: _Optional[str] = ..., target_variant: _Optional[str] = ..., images: _Optional[_Union[_base_pb2.TextArray, _Mapping]] = ..., image_properties: _Optional[_Union[_base_pb2.TextArray, _Mapping]] = ..., max_tokens: _Optional[int] = ...) -> None: ...
 
 class GenerativeCohere(_message.Message):
     __slots__ = ["base_url", "frequency_penalty", "max_tokens", "model", "k", "p", "presence_penalty", "stop_sequences", "temperature"]

weaviate/proto/v1/v4216/v1/generative_pb2.pyi

@@ -105,7 +105,7 @@ class GenerativeAnyscale(_message.Message):
     def __init__(self, base_url: _Optional[str] = ..., model: _Optional[str] = ..., temperature: _Optional[float] = ...) -> None: ...
 
 class GenerativeAWS(_message.Message):
-    __slots__ = ("model", "temperature", "service", "region", "endpoint", "target_model", "target_variant", "images", "image_properties")
+    __slots__ = ("model", "temperature", "service", "region", "endpoint", "target_model", "target_variant", "images", "image_properties", "max_tokens")
     MODEL_FIELD_NUMBER: _ClassVar[int]
     TEMPERATURE_FIELD_NUMBER: _ClassVar[int]
     SERVICE_FIELD_NUMBER: _ClassVar[int]
@@ -115,6 +115,7 @@ class GenerativeAWS(_message.Message):
     TARGET_VARIANT_FIELD_NUMBER: _ClassVar[int]
     IMAGES_FIELD_NUMBER: _ClassVar[int]
     IMAGE_PROPERTIES_FIELD_NUMBER: _ClassVar[int]
+    MAX_TOKENS_FIELD_NUMBER: _ClassVar[int]
     model: str
     temperature: float
     service: str
@@ -124,7 +125,8 @@ class GenerativeAWS(_message.Message):
     target_variant: str
     images: _base_pb2.TextArray
     image_properties: _base_pb2.TextArray
-    def __init__(self, model: _Optional[str] = ..., temperature: _Optional[float] = ..., service: _Optional[str] = ..., region: _Optional[str] = ..., endpoint: _Optional[str] = ..., target_model: _Optional[str] = ..., target_variant: _Optional[str] = ..., images: _Optional[_Union[_base_pb2.TextArray, _Mapping]] = ..., image_properties: _Optional[_Union[_base_pb2.TextArray, _Mapping]] = ...) -> None: ...
+    max_tokens: int
+    def __init__(self, model: _Optional[str] = ..., temperature: _Optional[float] = ..., service: _Optional[str] = ..., region: _Optional[str] = ..., endpoint: _Optional[str] = ..., target_model: _Optional[str] = ..., target_variant: _Optional[str] = ..., images: _Optional[_Union[_base_pb2.TextArray, _Mapping]] = ..., image_properties: _Optional[_Union[_base_pb2.TextArray, _Mapping]] = ..., max_tokens: _Optional[int] = ...) -> None: ...
 
 class GenerativeCohere(_message.Message):
     __slots__ = ("base_url", "frequency_penalty", "max_tokens", "model", "k", "p", "presence_penalty", "stop_sequences", "temperature")