Create Bedrock client

apps/openchallenges/notebook/notebooks/openai-challenge-headline.ipynb

@@ -253,12 +253,47 @@
     "print(json_str)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "14ba8e14",
+   "metadata": {},
+   "source": [
+    "## Generating challenge headlines with AWS LLM"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8ba9a632",
+   "metadata": {},
+   "source": [
+    "### Configure Bedrock client"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 19,
    "id": "20da8b0e",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "import json\n",
+    "import os\n",
+    "import sys\n",
+    "\n",
+    "import boto3\n",
+    "import botocore\n",
+    "\n",
+    "module_path = \"..\"\n",
+    "sys.path.append(os.path.abspath(module_path))\n",
+    "from utils import bedrock, print_ww"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5bbcd90f",
+   "metadata": {},
+   "outputs": [],
    "source": []
   }
  ],

apps/openchallenges/notebook/notebooks/utils/bedrock.py

@@ -0,0 +1,80 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+"""Helper utilities for working with Amazon Bedrock from Python notebooks"""
+# Python Built-Ins:
+import os
+from typing import Optional
+
+# External Dependencies:
+import boto3
+from botocore.config import Config
+
+
+def get_bedrock_client(
+    assumed_role: Optional[str] = None,
+    region: Optional[str] = None,
+    runtime: Optional[bool] = True,
+):
+    """Create a boto3 client for Amazon Bedrock, with optional configuration overrides
+
+    Parameters
+    ----------
+    assumed_role :
+        Optional ARN of an AWS IAM role to assume for calling the Bedrock service. If not
+        specified, the current active credentials will be used.
+    region :
+        Optional name of the AWS Region in which the service should be called (e.g. "us-east-1").
+        If not specified, AWS_REGION or AWS_DEFAULT_REGION environment variable will be used.
+    runtime :
+        Optional choice of getting different client to perform operations with the Amazon Bedrock service.
+    """
+    if region is None:
+        target_region = os.environ.get(
+            "AWS_REGION", os.environ.get("AWS_DEFAULT_REGION")
+        )
+    else:
+        target_region = region
+
+    print(f"Create new client\n  Using region: {target_region}")
+    session_kwargs = {"region_name": target_region}
+    client_kwargs = {**session_kwargs}
+
+    profile_name = os.environ.get("AWS_PROFILE")
+    if profile_name:
+        print(f"  Using profile: {profile_name}")
+        session_kwargs["profile_name"] = profile_name
+
+    retry_config = Config(
+        region_name=target_region,
+        retries={
+            "max_attempts": 10,
+            "mode": "standard",
+        },
+    )
+    session = boto3.Session(**session_kwargs)
+
+    if assumed_role:
+        print(f"  Using role: {assumed_role}", end="")
+        sts = session.client("sts")
+        response = sts.assume_role(
+            RoleArn=str(assumed_role), RoleSessionName="langchain-llm-1"
+        )
+        print(" ... successful!")
+        client_kwargs["aws_access_key_id"] = response["Credentials"]["AccessKeyId"]
+        client_kwargs["aws_secret_access_key"] = response["Credentials"][
+            "SecretAccessKey"
+        ]
+        client_kwargs["aws_session_token"] = response["Credentials"]["SessionToken"]
+
+    if runtime:
+        service_name = "bedrock-runtime"
+    else:
+        service_name = "bedrock"
+
+    bedrock_client = session.client(
+        service_name=service_name, config=retry_config, **client_kwargs
+    )
+
+    print("boto3 Bedrock client successfully created!")
+    print(bedrock_client._endpoint)
+    return bedrock_client

apps/openchallenges/notebook/notebooks/utils/print_ww.py

@@ -0,0 +1,21 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+"""General helper utilities the workshop notebooks"""
+# Python Built-Ins:
+from io import StringIO
+import sys
+import textwrap
+
+
+def print_ww(*args, width: int = 100, **kwargs):
+    """Like print(), but wraps output to `width` characters (default 100)"""
+    buffer = StringIO()
+    try:
+        _stdout = sys.stdout
+        sys.stdout = buffer
+        print(*args, **kwargs)
+        output = buffer.getvalue()
+    finally:
+        sys.stdout = _stdout
+    for line in output.splitlines():
+        print("\n".join(textwrap.wrap(line, width=width)))