bedrock_example.py

import os
import ldclient
from ldclient import Context
from ldclient.config import Config
from ldai import LDAIClient
import boto3

client = boto3.client("bedrock-runtime", region_name="us-east-1")

# Set sdk_key to your LaunchDarkly SDK key.
sdk_key = os.getenv('LAUNCHDARKLY_SDK_KEY')

# Set config_key to the AI Config key you want to evaluate.
ai_config_key = os.getenv('LAUNCHDARKLY_AI_CONFIG_KEY', 'sample-ai-config')

def main():
    if not sdk_key:
        print("*** Please set the LAUNCHDARKLY_SDK_KEY env first")
        exit()
    if not ai_config_key:
        print("*** Please set the LAUNCHDARKLY_AI_CONFIG_KEY env first")
        exit()

    ldclient.set_config(Config(sdk_key))

    if not ldclient.get().is_initialized():
        print("*** SDK failed to initialize. Please check your internet connection and SDK credential for any typo.")
        exit()

    aiclient = LDAIClient(ldclient.get())
    print("*** SDK successfully initialized")

    # Set up the evaluation context. This context should appear on your
    # LaunchDarkly contexts dashboard soon after you run the demo.
    context = (
        Context
        .builder('example-user-key')
        .kind('user')
        .name('Sandy')
        .build()
    )

    # Pass a default for improved resiliency when the AI config is unavailable
    # or LaunchDarkly is unreachable; omit for a disabled default.
    # Example:
    #   default = AIConfig(
    #       enabled=True,
    #       model=ModelConfig(name='my-default-model'),
    #       provider=ProviderConfig(name='bedrock'),
    #       messages=[LDMessage(role='system', content='You are a helpful assistant.')],
    #   )
    #   config_value = aiclient.completion_config(ai_config_key, context, default, {'myUserVariable': "Testing Variable"})
    config_value = aiclient.completion_config(
        ai_config_key,
        context,
        variables={'myUserVariable': "Testing Variable"}
    )
    tracker = config_value.tracker

    if not config_value.enabled:
        print("AI Config is disabled")
        return

    # Map the messages to the format expected by Bedrock
    chat_messages = [{'role': msg.role, 'content': [{'text': msg.content}]} for msg in config_value.messages if msg.role != 'system']
    system_messages = [{'text': msg.content} for msg in config_value.messages if msg.role == 'system']

    # Add the user input to the conversation
    USER_INPUT = "What can you help me with?"
    print("User Input:\n", USER_INPUT)
    chat_messages.append({'role': 'user', 'content': [{'text': USER_INPUT}]})

    converse = tracker.track_bedrock_converse_metrics(
        client.converse(
            modelId=config_value.model.name,
            messages=chat_messages,
            system=system_messages,
        )
    )

    # Append the AI response to the conversation history
    chat_messages.append(converse["output"]["message"])
    print("AI Response:\n", converse["output"]["message"]["content"][0]["text"])

    # Continue the conversation by adding user input to the messages list and invoking the LLM again.
    print("Success.")

    # Flush pending events and close the client.
    ldclient.get().flush()
    ldclient.get().close()