Skip to content

Add max token to AWS generative module #1819

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Sep 4, 2025
Merged

Add max token to AWS generative module #1819

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
fc57e55
Add max token to AWS generative module
dirkkul Sep 4, 2025
129b9a5
Make model a required parameter for text2vec-aws
dirkkul Sep 4, 2025
05bdc63
Fix test
dirkkul Sep 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion test/collection/test_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
],
{
Expand All @@ -2071,6 +2071,7 @@ def test_config_with_named_vectors(
"vectorizeClassName": True,
"region": "us-east-1",
"service": "bedrock",
"model": "model",
}
},
"vectorIndexType": "hnsw",
Expand Down
8 changes: 4 additions & 4 deletions weaviate/collections/classes/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -464,6 +464,7 @@ class _GenerativeAWSConfig(_GenerativeProvider):
service: str
model: Optional[str]
endpoint: Optional[str]
maxTokens: Optional[int]


class _GenerativeAnthropicConfig(_GenerativeProvider):
Expand Down Expand Up @@ -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.

Expand All @@ -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
Expand Down
4 changes: 2 additions & 2 deletions weaviate/collections/classes/config_vectors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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.
Expand Down
5 changes: 5 additions & 0 deletions weaviate/collections/classes/generative.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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),
),
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion weaviate/collections/classes/internal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
168 changes: 84 additions & 84 deletions weaviate/proto/v1/v4216/v1/generative_pb2.py
View file
Open in desktop

Large diffs are not rendered by default.

6 changes: 4 additions & 2 deletions weaviate/proto/v1/v4216/v1/generative_pb2.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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
Expand All @@ -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"]
Expand Down
168 changes: 84 additions & 84 deletions weaviate/proto/v1/v5261/v1/generative_pb2.py
View file
Open in desktop

Large diffs are not rendered by default.

6 changes: 4 additions & 2 deletions weaviate/proto/v1/v5261/v1/generative_pb2.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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
Expand All @@ -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")
Expand Down
168 changes: 84 additions & 84 deletions weaviate/proto/v1/v6300/v1/generative_pb2.py
View file
Open in desktop

Large diffs are not rendered by default.

6 changes: 4 additions & 2 deletions weaviate/proto/v1/v6300/v1/generative_pb2.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,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]
Expand All @@ -116,6 +116,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
Expand All @@ -125,7 +126,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")
Expand Down