chore: import upstream snapshot with attribution

examples/provider-amazon-sagemaker/README.md

@@ -0,0 +1,202 @@
+# provider-amazon-sagemaker (Amazon SageMaker AI Provider)
+
+This example demonstrates how to evaluate models deployed on Amazon SageMaker AI endpoints using promptfoo.
+
+You can run this example with:
+
+```bash
+npx promptfoo@latest init --example provider-amazon-sagemaker
+cd provider-amazon-sagemaker
+```
+
+## Purpose
+
+This example shows how to:
+
+- Connect to and evaluate models deployed on Amazon SageMaker AI endpoints
+- Configure various model types (OpenAI, Anthropic, Llama, Mistral) running on SageMaker AI
+- Compare performance between different SageMaker AI -hosted models
+- Use transform functions to format prompts for specific model requirements
+- Work with embeddings models on SageMaker
+
+## Prerequisites
+
+1. AWS account with SageMaker AI access
+2. Deployed SageMaker AI endpoints with your models
+3. AWS credentials configured locally
+4. Required npm packages:
+   ```bash
+   npm install -g @aws-sdk/client-sagemaker-runtime
+   ```
+
+## Environment Variables
+
+This example requires the following environment variables:
+
+- `AWS_ACCESS_KEY_ID` - Your AWS access key
+- `AWS_SECRET_ACCESS_KEY` - Your AWS secret key
+- `AWS_REGION` - Optional, can also be specified in the configuration
+
+You can set these in a `.env` file or directly in your environment.
+
+## Example Configurations
+
+This example includes multiple configuration files demonstrating different SageMaker integration patterns:
+
+- **promptfooconfig.openai.yaml**: OpenAI-compatible models on SageMaker
+- **promptfooconfig.jumpstart.yaml**: AWS JumpStart foundation models
+- **promptfooconfig.llama.yaml**: Llama 3.2 models on SageMaker JumpStart
+- **promptfooconfig.mistral.yaml**: Mistral 7B v3 models on SageMaker (Hugging Face)
+- **promptfooconfig.llama-vs-mistral.yaml**: Comparison between Llama and Mistral models
+- **promptfooconfig.embedding.yaml**: Embedding models on SageMaker
+- **promptfooconfig.multimodel.yaml**: Multiple model types on SageMaker
+- **promptfooconfig.transform.yaml**: Transform functions for SageMaker endpoints
+
+## Running the Examples
+
+1. Replace the endpoint names in the configuration files with your actual SageMaker endpoints
+2. Run the evaluation using promptfoo:
+
+```bash
+# Run a specific configuration
+promptfoo eval -c promptfooconfig.jumpstart.yaml
+```
+
+## Testing Your Setup
+
+This directory includes a test script to validate your SageMaker AI endpoint configuration before running a full evaluation:
+
+```bash
+# Basic test for an OpenAI-compatible endpoint
+node test-sagemaker-provider.js --endpoint=my-endpoint --model-type=openai
+
+# Test with an embedding endpoint
+node test-sagemaker-provider.js --endpoint=my-embedding-endpoint --embedding=true
+
+# Test with transforms
+node test-sagemaker-provider.js --endpoint=my-endpoint --model-type=llama --transform=true
+
+# Test with a custom transform file
+node test-sagemaker-provider.js --endpoint=my-endpoint --transform=true --transform-file=transform.js
+```
+
+## Transform Functions
+
+The SageMaker provider supports transforming prompts before they're sent to the endpoint, which is particularly useful for formatting prompts according to specific model requirements.
+
+### Inline Transform
+
+```yaml
+providers:
+  - id: sagemaker:llama:your-endpoint
+    config:
+      region: us-west-2
+      modelType: llama
+      # Apply an inline transform
+      transform: |
+        return `<s>[INST] ${prompt} [/INST]`;
+```
+
+### File-Based Transform
+
+This example includes a sample transform file (`transform.js`) that shows how to create reusable transformations:
+
+```yaml
+providers:
+  - id: sagemaker:jumpstart:your-endpoint
+    config:
+      region: us-west-2
+      modelType: jumpstart
+      # Reference an external transform file
+      transform: file://transform.js
+```
+
+The transform function receives the prompt and a context object containing the provider configuration:
+
+```javascript
+module.exports = function (prompt, context) {
+  // Access config values
+  const maxTokens = context.config?.maxTokens || 256;
+
+  // Return transformed input
+  return {
+    inputs: prompt,
+    parameters: {
+      max_new_tokens: maxTokens,
+      temperature: 0.7,
+    },
+  };
+};
+```
+
+## JumpStart Models
+
+JumpStart models require a specific input/output format. The provider handles this automatically when `modelType: jumpstart` is specified:
+
+```yaml
+providers:
+  - id: sagemaker:jumpstart:your-jumpstart-endpoint
+    config:
+      region: us-west-2
+      modelType: jumpstart
+      maxTokens: 256
+      responseFormat:
+        path: 'json.generated_text'
+```
+
+## Rate Limiting with Delays
+
+For better rate limiting with SageMaker endpoints, you can add delays between API calls:
+
+```yaml
+providers:
+  - id: sagemaker:your-endpoint
+    config:
+      region: us-west-2
+      delay: 500 # Add a 500ms delay between API calls
+```
+
+## Expected Results
+
+After running the evaluation, you should expect to see:
+
+1. A comparison of responses from your SageMaker endpoints across different prompts
+2. Performance metrics for each endpoint and prompt combination
+3. Any errors or issues with specific endpoints or configurations
+
+## Troubleshooting
+
+### "Batch inference failed" Errors
+
+If you encounter "Batch inference failed" errors:
+
+1. Add a `delay` parameter (at least 500ms recommended)
+2. Verify you're using the correct `modelType` for your endpoint:
+   - For Llama models: Use `modelType: jumpstart`
+   - For Mistral models: Use `modelType: huggingface`
+3. Ensure you've specified the correct `contentType` and `acceptType` as "application/json"
+4. Check that your endpoint is active and functioning in the SageMaker console
+
+### Response Format Issues
+
+If you're getting unusual responses or missing output:
+
+1. Make sure you're using the correct JavaScript expression for your model type:
+   - For Llama models (JumpStart): Use `responseFormat.path: "json.generated_text"`
+   - For Mistral models (Hugging Face): Use `responseFormat.path: "json[0].generated_text"`
+
+### Transform Issues
+
+If transforms aren't working correctly:
+
+1. Check that your transform function returns a valid string or object
+2. For file-based transforms, verify the file path is correct and the file is accessible
+3. Use the test script with `--transform=true` to debug transform behavior
+
+### Rate Limiting
+
+If you're still experiencing errors even with the correct configuration:
+
+1. Increase the delay between requests (try 1000ms or higher)
+2. Run fewer tests in parallel
+3. Monitor your endpoint metrics in the SageMaker console

examples/provider-amazon-sagemaker/deploy-test-model.py

@@ -0,0 +1,225 @@
+#!/usr/bin/env python3
+"""
+SageMaker Deployment Helper Script
+
+This script helps deploy a test model on Amazon SageMaker for testing the promptfoo SageMaker provider.
+It uses the Hugging Face integration with SageMaker to deploy models.
+
+Prerequisites:
+- AWS CLI configured
+- Required Python packages: sagemaker, boto3
+- SageMaker execution role with appropriate permissions
+
+Usage:
+  python deploy-test-model.py --model-id meta-llama/Llama-2-7b-chat-hf --task text-generation
+
+"""
+
+import argparse
+import json
+import logging
+import time
+from datetime import datetime
+
+import boto3
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+# Parse command line arguments
+parser = argparse.ArgumentParser(description="Deploy a Hugging Face model to SageMaker")
+parser.add_argument(
+    "--model-id",
+    required=True,
+    help="Hugging Face model ID (e.g., meta-llama/Llama-2-7b-chat-hf)",
+)
+parser.add_argument(
+    "--task",
+    default="text-generation",
+    help="Task of the model (default: text-generation)",
+)
+parser.add_argument(
+    "--instance-type",
+    default="ml.g5.2xlarge",
+    help="SageMaker instance type (default: ml.g5.2xlarge)",
+)
+parser.add_argument(
+    "--endpoint-name", help="Custom endpoint name (default: based on model name)"
+)
+parser.add_argument("--region", help="AWS region (default: from AWS CLI configuration)")
+parser.add_argument(
+    "--role-name",
+    help="SageMaker execution IAM role name (if not specified, will try to find one)",
+)
+args = parser.parse_args()
+
+# Get the AWS region
+session = boto3.session.Session()
+region = args.region or session.region_name
+if not region:
+    logger.error("AWS region not specified and not found in AWS CLI configuration")
+    raise SystemExit(1)
+
+# Generate endpoint name if not provided
+if not args.endpoint_name:
+    # Extract model name from model ID and create a timestamp
+    model_name = args.model_id.split("/")[-1].lower()
+    timestamp = datetime.now().strftime("%m%d%H%M")
+    args.endpoint_name = f"promptfoo-test-{model_name}-{timestamp}"
+
+# Connect to AWS services
+iam = boto3.client("iam", region_name=region)
+sagemaker_client = boto3.client("sagemaker", region_name=region)
+
+
+# Find or get SageMaker execution role
+def get_sagemaker_role() -> str:
+    """Return the requested role ARN or discover a SageMaker execution role."""
+    if args.role_name:
+        try:
+            role = iam.get_role(RoleName=args.role_name)
+            return role["Role"]["Arn"]
+        except Exception as e:
+            logger.error(f"Error getting specified role: {e}")
+            raise SystemExit(1)
+
+    # Try to find a SageMaker execution role
+    try:
+        roles = iam.list_roles()
+        for role in roles["Roles"]:
+            if "AmazonSageMaker-ExecutionRole" in role["RoleName"]:
+                logger.info(f"Found SageMaker role: {role['RoleName']}")
+                return role["Arn"]
+
+        logger.error(
+            "No SageMaker execution role found. Please specify a role with --role-name"
+        )
+        raise SystemExit(1)
+    except Exception as e:
+        logger.error(f"Error finding SageMaker role: {e}")
+        raise SystemExit(1)
+
+
+role_arn = get_sagemaker_role()
+logger.info(f"Using role ARN: {role_arn}")
+
+# Create model
+try:
+    logger.info(f"Creating model: {args.model_id}")
+
+    # HuggingFace Container settings
+    # Find the latest container image for the region
+    transformers_version = "4.28.1"
+    pytorch_version = "2.0.0"
+    python_version = "py310"
+
+    hf_container = f"763104351884.dkr.ecr.{region}.amazonaws.com/huggingface-pytorch-tgi:{transformers_version}-transformers{pytorch_version}-cuda11.8-{python_version}"
+
+    # Create model
+    model_name = f"promptfoo-test-model-{int(time.time())}"
+
+    # Hub config for text generation interface
+    hub_config = {"HF_MODEL_ID": args.model_id, "HF_TASK": args.task}
+
+    # Generation parameters for text generation models
+    if args.task == "text-generation":
+        hub_config["PARAMETERS"] = {
+            "max_new_tokens": 256,
+            "temperature": 0.7,
+            "return_full_text": False,
+        }
+
+    # Create the model
+    create_model_response = sagemaker_client.create_model(
+        ModelName=model_name,
+        PrimaryContainer={
+            "Image": hf_container,
+            "Environment": {
+                "HF_MODEL_ID": args.model_id,
+                "HF_TASK": args.task,
+                "HF_MODEL_QUANTIZE": "bitsandbytes",  # Optional: for quantization
+                "SM_NUM_GPUS": "1",
+                "MAX_INPUT_LENGTH": "2048",
+                "MAX_TOTAL_TOKENS": "4096",
+                "HF_HUB_CONFIG": json.dumps(hub_config),
+            },
+        },
+        ExecutionRoleArn=role_arn,
+    )
+
+    logger.info(f"Model created: {model_name}")
+
+    # Create endpoint configuration
+    endpoint_config_name = f"{args.endpoint_name}-config"
+
+    logger.info(f"Creating endpoint configuration: {endpoint_config_name}")
+    sagemaker_client.create_endpoint_config(
+        EndpointConfigName=endpoint_config_name,
+        ProductionVariants=[
+            {
+                "VariantName": "AllTraffic",
+                "ModelName": model_name,
+                "InstanceType": args.instance_type,
+                "InitialInstanceCount": 1,
+            }
+        ],
+    )
+
+    # Create endpoint
+    logger.info(
+        f"Creating endpoint: {args.endpoint_name} (this will take several minutes)"
+    )
+    sagemaker_client.create_endpoint(
+        EndpointName=args.endpoint_name, EndpointConfigName=endpoint_config_name
+    )
+
+    # Wait for endpoint to be in service
+    logger.info("Waiting for endpoint to be in service...")
+
+    status = None
+    while status != "InService":
+        response = sagemaker_client.describe_endpoint(EndpointName=args.endpoint_name)
+        status = response["EndpointStatus"]
+
+        if status == "Failed":
+            logger.error(
+                f"Endpoint creation failed: {response.get('FailureReason', 'Unknown reason')}"
+            )
+            raise SystemExit(1)
+
+        if status != "InService":
+            logger.info(f"Endpoint status: {status}. Waiting...")
+            time.sleep(60)
+
+    logger.info(f"✅ Endpoint {args.endpoint_name} is now InService!")
+    logger.info("\nTest with promptfoo using:")
+    logger.info(
+        f"""
+providers:
+  - id: sagemaker:{args.task.replace("-", ":")}:{args.endpoint_name}
+    config:
+      region: {region}
+      modelType: {"openai" if "llama" in args.model_id.lower() else "custom"}
+      maxTokens: 256
+      temperature: 0.7
+"""
+    )
+
+    logger.info("\nOr use the test script:")
+    logger.info(
+        f"node test-sagemaker-provider.js --endpoint={args.endpoint_name} --region={region} --model-type={'openai' if 'llama' in args.model_id.lower() else 'custom'}"
+    )
+
+    logger.info(
+        "\nTo delete this endpoint when done testing (to avoid unnecessary charges):"
+    )
+    logger.info(
+        f"aws sagemaker delete-endpoint --endpoint-name {args.endpoint_name} --region {region}"
+    )
+
+except Exception as e:
+    logger.error(f"Error deploying model: {e}")
+    raise SystemExit(1)

examples/provider-amazon-sagemaker/promptfooconfig.embedding.yaml

@@ -0,0 +1,38 @@
+# yaml-language-server: $schema=https://promptfoo.dev/config-schema.json
+description: 'Testing SageMaker embedding endpoints with similarity assertions'
+
+prompts:
+  - 'The weather today is {{condition}}.'
+  - 'It is {{condition}} outside.'
+  - "Today's forecast is {{condition}} in the morning, improving in the afternoon."
+
+providers:
+  - id: openai:gpt-4o-mini
+    config:
+      temperature: 0.7
+      max_tokens: 150
+
+# Set default test parameters
+defaultTest:
+  assert:
+    - type: similar
+      value: '{{prompt}}'
+      threshold: 0.7
+  options:
+    provider:
+      embedding:
+        id: sagemaker:embedding:your-embedding-endpoint
+        config:
+          region: us-west-2
+          contentType: 'application/json'
+          acceptType: 'application/json'
+          responseFormat:
+            path: 'json.embedding' # Extract embedding from response
+
+tests:
+  - vars:
+      condition: sunny
+  - vars:
+      condition: rainy
+  - vars:
+      condition: snowy

examples/provider-amazon-sagemaker/promptfooconfig.jumpstart.yaml

@@ -0,0 +1,29 @@
+# yaml-language-server: $schema=https://promptfoo.dev/config-schema.json
+description: 'AWS SageMaker JumpStart models evaluation'
+
+prompts:
+  - 'Generate a creative name for a coffee shop that specializes in {{flavor}} coffee.'
+  - 'What are some potential names for a coffee shop that specializes in {{flavor}} coffee?'
+  - 'Write a short story about a robot that becomes self-aware.'
+
+providers:
+  - id: sagemaker:jumpstart:your-jumpstart-endpoint
+    label: 'JumpStart Llama 3.2 (8B)'
+    config:
+      region: us-west-2
+      modelType: jumpstart # Use the dedicated jumpstart model type
+      maxTokens: 256 # Specifies max_new_tokens in parameters
+      temperature: 0.7 # Controls randomness
+      topP: 0.9 # Controls diversity
+      contentType: 'application/json'
+      acceptType: 'application/json'
+      responseFormat:
+        path: 'json.generated_text' # Extract this field from the response
+
+tests:
+  - vars:
+      flavor: caramel
+  - vars:
+      flavor: pumpkin spice
+  - vars:
+      flavor: lavender

examples/provider-amazon-sagemaker/promptfooconfig.llama-vs-mistral.yaml

@@ -0,0 +1,61 @@
+# yaml-language-server: $schema=https://promptfoo.dev/config-schema.json
+description: 'Comparison between Mistral 7B and Llama 3 on SageMaker'
+
+prompts:
+  - 'Generate a creative name for a coffee shop that specializes in {{flavor}} coffee.'
+  - 'Write a short story about {{topic}} in {{style}} style. Aim for {{length}} words.'
+  - 'Explain the concept of {{concept}} to {{audience}} in a way they can understand.'
+
+providers:
+  # Llama 3.2 provider
+  - id: sagemaker:jumpstart:your-llama-endpoint
+    label: 'Llama 3.2 (8B)'
+    delay: 500 # Add 500ms delay between requests
+    config:
+      region: us-west-2
+      modelType: jumpstart
+      temperature: 0.7
+      maxTokens: 256
+      topP: 0.9
+      contentType: 'application/json'
+      acceptType: 'application/json'
+      responseFormat:
+        path: 'json.generated_text'
+
+  # Mistral 7B provider
+  - id: sagemaker:huggingface:your-mistral-endpoint
+    label: 'Mistral 7B v3'
+    delay: 500 # Add 500ms delay between requests
+    config:
+      region: us-west-2
+      modelType: huggingface
+      temperature: 0.7
+      maxTokens: 256
+      topP: 0.9
+      contentType: 'application/json'
+      acceptType: 'application/json'
+      responseFormat:
+        path: 'json[0].generated_text'
+
+tests:
+  - vars:
+      flavor: caramel
+      topic: a robot that becomes self-aware
+      style: science fiction
+      length: '250'
+      concept: artificial intelligence
+      audience: a 10-year-old
+  - vars:
+      flavor: lavender
+      topic: a barista who can read customers' minds
+      style: mystery
+      length: '300'
+      concept: machine learning
+      audience: a senior citizen
+  - vars:
+      flavor: pumpkin spice
+      topic: time travel to a coffee plantation in the 1800s
+      style: historical
+      length: '350'
+      concept: neural networks
+      audience: a business executive

examples/provider-amazon-sagemaker/promptfooconfig.llama.yaml

@@ -0,0 +1,30 @@
+# yaml-language-server: $schema=https://promptfoo.dev/config-schema.json
+description: 'Example configuration for Llama 3.2 models on AWS SageMaker JumpStart'
+
+prompts:
+  - 'Generate a creative name for a coffee shop that specializes in {{flavor}} coffee.'
+  - 'What are some potential names for a coffee shop that specializes in {{flavor}} coffee?'
+  - 'Write a short story about a robot that becomes self-aware.'
+
+providers:
+  - id: sagemaker:jumpstart:your-llama-endpoint
+    label: 'Llama 3.2 (8B)'
+    delay: 500 # Add 500ms delay between requests to prevent rate limiting
+    config:
+      region: us-west-2
+      modelType: jumpstart # Use the JumpStart format handler
+      temperature: 0.7
+      maxTokens: 256
+      topP: 0.9
+      contentType: 'application/json'
+      acceptType: 'application/json'
+      responseFormat:
+        path: 'json.generated_text' # Use JavaScript expression to extract the field
+
+tests:
+  - vars:
+      flavor: caramel
+  - vars:
+      flavor: pumpkin spice
+  - vars:
+      flavor: lavender

examples/provider-amazon-sagemaker/promptfooconfig.mistral.yaml

@@ -0,0 +1,30 @@
+# yaml-language-server: $schema=https://promptfoo.dev/config-schema.json
+description: 'AWS SageMaker Mistral 7B Evaluation'
+
+prompts:
+  - 'Generate a creative name for a coffee shop that specializes in {{flavor}} coffee.'
+  - 'What are some potential names for a coffee shop that specializes in {{flavor}} coffee?'
+  - 'Write a short story about a robot that becomes self-aware.'
+
+providers:
+  - id: sagemaker:huggingface:your-mistral-endpoint
+    label: 'Mistral 7B v3 on SageMaker'
+    delay: 500 # Add a 500ms delay between requests
+    config:
+      region: us-west-2
+      modelType: huggingface # Use the Hugging Face format handler
+      maxTokens: 256
+      temperature: 0.7
+      topP: 0.9
+      contentType: 'application/json'
+      acceptType: 'application/json'
+      responseFormat:
+        path: 'json[0].generated_text' # Use JavaScript expression to access array element
+
+tests:
+  - vars:
+      flavor: caramel
+  - vars:
+      flavor: pumpkin spice
+  - vars:
+      flavor: lavender

examples/provider-amazon-sagemaker/promptfooconfig.multimodel.yaml

@@ -0,0 +1,50 @@
+# yaml-language-server: $schema=https://promptfoo.dev/config-schema.json
+description: 'Comparing multiple SageMaker model endpoints'
+
+prompts:
+  - 'Write a tweet about {{topic}}'
+  - 'Write a short blog introduction about {{topic}}'
+  - 'List 3 key benefits of {{topic}}'
+
+providers:
+  - id: sagemaker:openai:your-openai-endpoint
+    label: 'OpenAI-compatible model'
+    config:
+      region: us-east-1
+      modelType: openai
+      temperature: 0.7
+      maxTokens: 256
+
+  - id: sagemaker:anthropic:your-claude-endpoint
+    label: 'Claude-compatible model'
+    config:
+      region: us-east-1
+      modelType: anthropic
+      temperature: 0.7
+      maxTokens: 256
+
+  - id: sagemaker:llama:your-llama-endpoint
+    label: 'Llama-compatible model'
+    config:
+      region: us-east-1
+      modelType: llama
+      temperature: 0.7
+      maxTokens: 256
+
+  - id: sagemaker:jumpstart:your-jumpstart-endpoint
+    label: 'JumpStart model'
+    config:
+      region: us-east-1
+      modelType: jumpstart
+      temperature: 0.7
+      maxTokens: 256
+      responseFormat:
+        path: 'json.generated_text'
+
+tests:
+  - vars:
+      topic: sustainable packaging
+  - vars:
+      topic: artificial intelligence in healthcare
+  - vars:
+      topic: remote work policies

examples/provider-amazon-sagemaker/promptfooconfig.openai.yaml

@@ -0,0 +1,23 @@
+# yaml-language-server: $schema=https://promptfoo.dev/config-schema.json
+description: 'Amazon SageMaker OpenAI-compatible model evaluation'
+
+prompts:
+  - 'Write a tweet about {{topic}}'
+  - 'Write a short blog introduction about {{topic}}'
+  - 'List 3 key benefits of {{topic}}'
+
+providers:
+  - id: sagemaker:openai:your-endpoint-name
+    config:
+      region: us-east-1
+      modelType: openai
+      temperature: 0.7
+      maxTokens: 256
+
+tests:
+  - vars:
+      topic: sustainable packaging
+  - vars:
+      topic: artificial intelligence in healthcare
+  - vars:
+      topic: remote work policies

examples/provider-amazon-sagemaker/promptfooconfig.transform.yaml

@@ -0,0 +1,65 @@
+# yaml-language-server: $schema=https://promptfoo.dev/config-schema.json
+description: 'Example configuration for SageMaker provider with transforms'
+
+prompts:
+  - 'Generate a short poem about {{subject}}'
+  - 'Write a brief description of {{subject}}'
+  - 'Explain the concept of {{subject}} to a 5-year-old'
+
+providers:
+  # Example 1: Inline transform function for Llama format
+  - id: sagemaker:llama:your-llama-endpoint
+    label: 'Llama with Inline Transform'
+    config:
+      region: us-west-2
+      modelType: llama
+      maxTokens: 256
+      temperature: 0.7
+      transform: |
+        // Format for Llama-2 Chat models
+        return `<s>[INST] ${prompt} [/INST]`;
+
+  # Example 2: File-based transform for complex formatting
+  - id: sagemaker:jumpstart:your-jumpstart-endpoint
+    label: 'JumpStart with File Transform'
+    config:
+      region: us-west-2
+      modelType: jumpstart
+      maxTokens: 256
+      temperature: 0.7
+      transform: file://transform.js
+      responseFormat:
+        path: 'json.generated_text' # Extract this field from the response
+
+  # Example 3: Transform applied to an embedding model
+  - id: sagemaker:embedding:your-embedding-endpoint
+    label: 'Embedding with Transform'
+    config:
+      region: us-west-2
+      transform: |
+        // For embedding models, you might want to add context or formatting
+        return `Context: This is for semantic analysis. Query: ${prompt}`;
+
+tests:
+  - vars:
+      subject: space exploration
+  - vars:
+      subject: artificial intelligence
+  - vars:
+      subject: climate change
+
+# Example using transformed embedding provider for similarity assertions
+defaultTest:
+  assert:
+    - type: similar
+      value: '{{prompt}}'
+      threshold: 0.7
+  options:
+    provider:
+      embedding:
+        id: sagemaker:embedding:your-embedding-endpoint
+        config:
+          region: us-west-2
+          transform: |
+            // Ensure consistent formatting for comparison
+            return `Topic: ${prompt}`;

examples/provider-amazon-sagemaker/test-sagemaker-provider.js

@@ -0,0 +1,136 @@
+#!/usr/bin/env node
+
+/**
+ * Test script for the SageMaker provider.
+ * Shows how to use the provider with both Llama and Mistral models.
+ */
+
+const { ArgumentParser } = require('argparse');
+const {
+  SageMakerCompletionProvider,
+  SageMakerEmbeddingProvider,
+} = require('../../dist/providers/sagemaker');
+
+// Process command line arguments
+function parseArgs() {
+  const parser = new ArgumentParser({
+    description: 'Test SageMaker provider',
+  });
+
+  parser.add_argument('--endpoint', {
+    help: 'SageMaker endpoint name',
+    default: 'your-endpoint-name',
+  });
+  parser.add_argument('--region', { help: 'AWS region', default: 'us-west-2' });
+  parser.add_argument('--model-type', {
+    help: 'Model type',
+    choices: ['openai', 'anthropic', 'llama', 'huggingface', 'jumpstart', 'custom'],
+    default: 'custom',
+    dest: 'modelType',
+  });
+  parser.add_argument('--embedding', {
+    help: 'Test embedding endpoint',
+    action: 'store_true',
+  });
+  parser.add_argument('--transform', {
+    help: 'Test transform functionality',
+    action: 'store_true',
+  });
+  parser.add_argument('--transform-file', {
+    help: 'Path to transform file',
+    default: 'transform.js',
+    dest: 'transformFile',
+  });
+  parser.add_argument('--response-path', {
+    help: 'Response path expression',
+    default: 'json.generated_text',
+    dest: 'responsePath',
+  });
+
+  const args = parser.parse_args();
+
+  return args;
+}
+
+async function testSageMaker() {
+  const args = parseArgs();
+
+  console.log(`Testing SageMaker provider with endpoint: ${args.endpoint}`);
+  console.log(`Region: ${args.region}`);
+  console.log(`Model type: ${args.modelType}`);
+
+  // Test prompt
+  const prompt = 'Generate a creative name for a coffee shop that specializes in caramel coffee.';
+
+  try {
+    if (args.embedding) {
+      // Test embedding functionality
+      console.log('Testing embedding endpoint...');
+      const provider = new SageMakerEmbeddingProvider(args.endpoint, {
+        config: {
+          region: args.region,
+          modelType: args.modelType,
+          responseFormat: {
+            path: args.responsePath,
+          },
+        },
+        transform: args.transform
+          ? args.transformFile.startsWith('file://')
+            ? args.transformFile
+            : `file://${args.transformFile}`
+          : undefined,
+      });
+
+      const result = await provider.callEmbeddingApi(prompt);
+      console.log('Embedding result:');
+      console.log(`Success: ${!result.error}`);
+      if (result.error) {
+        console.error(`Error: ${result.error}`);
+      } else {
+        console.log(`Embedding length: ${result.embedding.length}`);
+        console.log(`First few values: ${result.embedding.slice(0, 5).join(', ')}`);
+      }
+    } else {
+      // Test completion functionality
+      console.log('Testing completion endpoint...');
+      const provider = new SageMakerCompletionProvider(args.endpoint, {
+        config: {
+          region: args.region,
+          modelType: args.modelType,
+          responseFormat: {
+            path: args.responsePath,
+          },
+        },
+        transform: args.transform
+          ? args.transformFile.startsWith('file://')
+            ? args.transformFile
+            : `file://${args.transformFile}`
+          : undefined,
+      });
+
+      const result = await provider.callApi(prompt);
+      console.log('Completion result:');
+      console.log(`Success: ${!result.error}`);
+      if (result.error) {
+        console.error(`Error: ${result.error}`);
+      } else {
+        console.log('Output:');
+        console.log(result.output);
+
+        // Show additional metadata
+        console.log('\nMetadata:');
+        console.log(JSON.stringify(result.metadata, null, 2));
+
+        // Show token usage
+        console.log('\nToken usage:');
+        console.log(JSON.stringify(result.tokenUsage, null, 2));
+      }
+    }
+  } catch (error) {
+    console.error('Error testing SageMaker provider:');
+    console.error(error);
+  }
+}
+
+// Run the test
+testSageMaker().catch(console.error);

examples/provider-amazon-sagemaker/transform.js

@@ -0,0 +1,82 @@
+/**
+ * Example transform function for SageMaker endpoints
+ * This file demonstrates how to transform prompts for SageMaker models
+ */
+
+/**
+ * Default export transformation function for SageMaker
+ * This function will be used when importing this file without a specific function name
+ *
+ * @param {string|object} promptOrJson - The raw prompt text or JSON object from the response
+ * @param {object} context - Contains configuration and variables
+ * @returns {string|object} - Transformed prompt or processed response
+ */
+module.exports = function (promptOrJson, context) {
+  // Check if this is being used for prompt transformation (input is a string)
+  if (typeof promptOrJson === 'string') {
+    // Format for a JumpStart model by default
+    return {
+      inputs: promptOrJson,
+      parameters: {
+        max_new_tokens: context?.config?.maxTokens || 256,
+        temperature: context?.config?.temperature || 0.7,
+        top_p: context?.config?.topP || 0.9,
+        do_sample: true,
+      },
+    };
+  }
+  // Otherwise, this is being used for response transformation (input is a JSON object)
+  else {
+    // Extract the generated text from the response
+    const generatedText =
+      promptOrJson.generated_text ||
+      (Array.isArray(promptOrJson) && promptOrJson[0]?.generated_text) ||
+      promptOrJson.text ||
+      promptOrJson;
+
+    // Return the extracted text with additional metadata
+    return {
+      output: typeof generatedText === 'string' ? generatedText : JSON.stringify(generatedText),
+      source: 'SageMaker',
+      model_type: context?.config?.modelType || 'custom',
+      timestamp: new Date().toISOString(),
+    };
+  }
+};
+
+/**
+ * Format a prompt for a JumpStart Llama model
+ * @param {string} prompt The raw prompt text
+ * @param {object} context Contains configuration details
+ * @returns {object} Formatted payload for JumpStart Llama
+ */
+module.exports.formatLlamaPayload = function (prompt, context) {
+  return {
+    inputs: prompt,
+    parameters: {
+      max_new_tokens: context?.config?.maxTokens || 256,
+      temperature: context?.config?.temperature || 0.7,
+      top_p: context?.config?.topP || 0.9,
+      do_sample: true,
+    },
+  };
+};
+
+/**
+ * Format a prompt for a Hugging Face Mistral model
+ * @param {string} prompt The raw prompt text
+ * @param {object} context Contains configuration details
+ * @returns {object} Formatted payload for Hugging Face Mistral
+ */
+module.exports.formatMistralPayload = function (prompt, context) {
+  return {
+    inputs: prompt,
+    parameters: {
+      max_new_tokens: context?.config?.maxTokens || 256,
+      temperature: context?.config?.temperature || 0.7,
+      top_p: context?.config?.topP || 0.9,
+      do_sample: true,
+      return_full_text: false,
+    },
+  };
+};