SagemakerGenerator

Generate text using large language models deployed on Amazon SageMaker Inference Endpoints.

Key Features

• Works with any LLM deployed on an Amazon SageMaker Inference Endpoint.
• Supports SageMaker JumpStart foundation models including Llama, Falcon, and others.
• Accepts a text prompt and returns generated text responses.
• AWS credentials are resolved through Secret parameters or the standard boto3 credential chain.
• Configurable generation parameters through generation_kwargs.

Configuration

1. Drag the SagemakerGenerator component onto the canvas from the Component Library.
2. Click on the component to open the configuration panel.
3. On the General tab:
   1. Set the model to the name of your SageMaker Inference Endpoint.
   2. Create secrets for your AWS credentials. Use AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, and AWS_DEFAULT_REGION as the environment variable names. For instructions, see Create Secrets.
4. Go to the Advanced tab to configure generation_kwargs (for example, max_new_tokens) and aws_custom_attributes.

Connections

SagemakerGenerator receives a text prompt string. It outputs a list of text replies and associated metadata.

Source Code

To check this component's source code, open sagemaker.py in the Haystack Core Integrations repository.

Usage Examples

Basic Configuration

  SagemakerGenerator:
    type: haystack_integrations.components.generators.amazon_sagemaker.SagemakerGenerator
    init_parameters:
      model: my-sagemaker-endpoint-name
      aws_access_key_id:
        type: env_var
        env_vars:
        - AWS_ACCESS_KEY_ID
        strict: false
      aws_secret_access_key:
        type: env_var
        env_vars:
        - AWS_SECRET_ACCESS_KEY
        strict: false
      aws_region_name:
        type: env_var
        env_vars:
        - AWS_DEFAULT_REGION
        strict: false
      generation_kwargs:
        max_new_tokens: 1024

Using the Component in a Pipeline

# haystack-pipeline
components:
  prompt_builder:
    type: haystack.components.builders.prompt_builder.PromptBuilder
    init_parameters:
      required_variables: "*"
      template: "Answer this question: {{ question }}"

  llm:
    type: haystack_integrations.components.generators.amazon_sagemaker.SagemakerGenerator
    init_parameters:
      model: my-sagemaker-endpoint-name
      aws_access_key_id:
        type: env_var
        env_vars:
        - AWS_ACCESS_KEY_ID
        strict: false
      aws_secret_access_key:
        type: env_var
        env_vars:
        - AWS_SECRET_ACCESS_KEY
        strict: false
      aws_region_name:
        type: env_var
        env_vars:
        - AWS_DEFAULT_REGION
        strict: false
      generation_kwargs:
        max_new_tokens: 512

connections:
- sender: prompt_builder.prompt
  receiver: llm.prompt

max_runs_per_component: 100

metadata: {}

inputs:
  query:
  - prompt_builder.question

Parameters

Inputs

Parameter	Type	Description
prompt	str	The text prompt to send to the model.
generation_kwargs	Optional[Dict[str, Any]]	Generation parameters to override init-time values.

Outputs

Parameter	Type	Description
replies	List[str]	A list of generated text responses.
meta	List[Dict[str, Any]]	Metadata about each generation response.

Init Parameters

These are the parameters you can configure in Pipeline Builder:

Parameter	Type	Default	Description
model	str	(required)	The name of the SageMaker Inference Endpoint to call.
aws_access_key_id	Optional[Secret]	Secret.from_env_var("AWS_ACCESS_KEY_ID", strict=False)	The AWS access key ID.
aws_secret_access_key	Optional[Secret]	Secret.from_env_var("AWS_SECRET_ACCESS_KEY", strict=False)	The AWS secret access key.
aws_session_token	Optional[Secret]	Secret.from_env_var("AWS_SESSION_TOKEN", strict=False)	The AWS session token for temporary credentials.
aws_region_name	Optional[Secret]	Secret.from_env_var("AWS_DEFAULT_REGION", strict=False)	The AWS region where the endpoint is deployed.
aws_profile_name	Optional[Secret]	Secret.from_env_var("AWS_PROFILE", strict=False)	The AWS named profile to use.
aws_custom_attributes	Optional[Dict[str, Any]]	None	Custom attributes to pass to the SageMaker endpoint invocation.
generation_kwargs	Optional[Dict[str, Any]]	{"max_new_tokens": 1024}	Default generation parameters sent to the SageMaker endpoint, such as max_new_tokens, temperature, and top_p.

Run Method Parameters

These are the parameters you can configure for the component's run() method. This means you can pass these parameters at query time through the API, in Playground, or when running a job. For details, see Modify Pipeline Parameters at Query Time.

Parameter	Type	Default	Description
prompt	str		The text prompt to send to the model.
generation_kwargs	Optional[Dict[str, Any]]	None	Generation parameters to override init-time values.

Related Information

• Pipeline Components