Observability for AWS Bedrock with Opik

1
import boto3
2
3
REGION = "us-east-1"
4
MODEL_ID = "us.meta.llama3-2-3b-instruct-v1:0"
5
6
bedrock = boto3.client(
7
    service_name="bedrock-runtime",
8
    region_name=REGION,
9
    # aws_access_key_id=ACCESS_KEY,
10
    # aws_secret_access_key=SECRET_KEY,
11
    # aws_session_token=SESSION_TOKEN,
12
)

1
from opik.integrations.bedrock import track_bedrock
2
3
bedrock_client = track_bedrock(bedrock, project_name="bedrock-integration-demo")
4
5
PROMPT = "Why is it important to use a LLM Monitoring like CometML Opik tool that allows you to log traces and spans when working with LLM Models hosted on AWS Bedrock?"
6
7
response = bedrock_client.converse(
8
    modelId=MODEL_ID,
9
    messages=[{"role": "user", "content": [{"text": PROMPT}]}],
10
    inferenceConfig={"temperature": 0.5, "maxTokens": 512, "topP": 0.9},
11
)
12
print("Response", response["output"]["message"]["content"][0]["text"])

1
def stream_conversation(
2
    bedrock_client,
3
    model_id,
4
    messages,
5
    system_prompts,
6
    inference_config,
7
):
8
    """
9
    Sends messages to a model and streams the response.
10
    Args:
11
        bedrock_client: The Boto3 Bedrock runtime client.
12
        model_id (str): The model ID to use.
13
        messages (JSON) : The messages to send.
14
        system_prompts (JSON) : The system prompts to send.
15
        inference_config (JSON) : The inference configuration to use.
16
17
    Returns:
18
        Nothing.
19
    """
20
21
    response = bedrock_client.converse_stream(
22
        modelId=model_id,
23
        messages=messages,
24
        system=system_prompts,
25
        inferenceConfig=inference_config,
26
    )
27
28
    stream = response.get("stream")
29
    if stream:
30
        for event in stream:
31
            if "messageStart" in event:
32
                print(f"\nRole: {event['messageStart']['role']}")
33
34
            if "contentBlockDelta" in event:
35
                print(event["contentBlockDelta"]["delta"]["text"], end="")
36
37
            if "messageStop" in event:
38
                print(f"\nStop reason: {event['messageStop']['stopReason']}")
39
40
            if "metadata" in event:
41
                metadata = event["metadata"]
42
                if "usage" in metadata:
43
                    print("\nToken usage")
44
                    print(f"Input tokens: {metadata['usage']['inputTokens']}")
45
                    print(f"Output tokens: {metadata['usage']['outputTokens']}")
46
                    print(f"Total tokens: {metadata['usage']['totalTokens']}")
47
                if "metrics" in event["metadata"]:
48
                    print(f"Latency: {metadata['metrics']['latencyMs']} milliseconds")
49
50
# Example usage with system prompts
51
system_prompt = """You are an app that creates playlists for a radio station
52
  that plays rock and pop music. Only return song names and the artist."""
53
54
input_text = "Create a list of 3 pop songs."
55
56
message = {"role": "user", "content": [{"text": input_text}]}
57
messages = [message]
58
59
# System prompts.
60
system_prompts = [{"text": system_prompt}]
61
62
# inference parameters to use.
63
inference_config = {"temperature": 0.5, "topP": 0.9}
64
65
stream_conversation(
66
    bedrock_client,
67
    MODEL_ID,
68
    messages,
69
    system_prompts,
70
    inference_config,
71
)

1
from opik import track
2
from opik.integrations.bedrock import track_bedrock
3
4
os.environ["OPIK_PROJECT_NAME"] = "bedrock-integration-demo"
5
bedrock_client = track_bedrock(bedrock)
6
7
@track
8
def generate_story(prompt):
9
    res = bedrock_client.converse(
10
        modelId=MODEL_ID, messages=[{"role": "user", "content": [{"text": prompt}]}]
11
    )
12
    return res["output"]["message"]["content"][0]["text"]
13
14
@track
15
def generate_topic():
16
    prompt = "Generate a topic for a story about Opik."
17
    res = bedrock_client.converse(
18
        modelId=MODEL_ID, messages=[{"role": "user", "content": [{"text": prompt}]}]
19
    )
20
    return res["output"]["message"]["content"][0]["text"]
21
22
@track
23
def generate_opik_story():
24
    topic = generate_topic()
25
    story = generate_story(topic)
26
    return story
27
28
# Execute the multi-step pipeline
29
generate_opik_story()