Skip to content

Commit 957d45c

Browse files
Fix: [AEA-6112] - kb logging 🤞 (#271)
## summary adds configurable aws bedrock model invocation logging so we can capture + monitor all bedrock api interactions in cloudwatch logs ### details this pr adds the infra + lambda glue to enable aws bedrock model invocation logging cdk doesn’t currently expose bedrock’s logging config, so we wire it up via a custom cloudformation resource (backed by a lambda) that calls the bedrock api logging is **off by default**. you can toggle it via the `enableLogging` flag in `EpsAssistMeStack.ts` during deployment or by invoking the lambda in aws console with the following body: ``` {"enable_logging": true} OR {"enable_logging": false} ``` - **BedrockLoggingConfiguration construct** - creates a cloudwatch log group (kms encrypted) for invocation logs - creates an iam role that the bedrock service can assume to write logs - creates a custom resource lambda to apply/remove the bedrock logging config - **bedrockLoggingConfigFunction lambda** - on create/update: applies the model invocation logging config - on delete: removes the logging config - supports toggling via `ENABLE_LOGGING` env var (so deploys can keep the resource but no-op when disabled) - derives log group + role arns from cdk-provided env vars --------- Co-authored-by: Beenyaa <bencegadanyi1@hotmail.com>
1 parent eb14a27 commit 957d45c

15 files changed

Lines changed: 962 additions & 124 deletions

‎.github/scripts/fix_cdk_json.sh‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ fix_string_key versionNumber "${VERSION_NUMBER}"
5757
fix_string_key commitId "${COMMIT_ID}"
5858
fix_string_key logRetentionInDays "${LOG_RETENTION_IN_DAYS}"
5959
fix_string_key logLevel "${LOG_LEVEL}"
60+
fix_string_key enableBedrockLogging "${ENABLE_BEDROCK_LOGGING:-false}"
6061
fix_string_key slackBotToken "${SLACK_BOT_TOKEN}"
6162
fix_string_key slackSigningSecret "${SLACK_SIGNING_SECRET}"
6263
fix_string_key cfnDriftDetectionGroup "${CFN_DRIFT_DETECTION_GROUP}"

‎.github/workflows/cdk_package_code.yml‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ jobs:
6868
run: |
6969
poetry show --only=slackBotFunction | grep -E "^[a-zA-Z]" | awk '{print $1"=="$2}' > requirements_slackBotFunction
7070
poetry show --only=syncKnowledgeBaseFunction | grep -E "^[a-zA-Z]" | awk '{print $1"=="$2}' > requirements_syncKnowledgeBaseFunction
71+
poetry show --only=bedrockLoggingConfigFunction | grep -E "^[a-zA-Z]" | awk '{print $1"=="$2}' > requirements_bedrockLoggingConfigFunction
7172
if [ ! -s requirements_slackBotFunction ] || [ "$(grep -c -v '^[[:space:]]*$' requirements_slackBotFunction)" -eq 0 ]; then \
7273
echo "Error: requirements_slackBotFunction is empty or contains only blank lines"; \
7374
exit 1; \
@@ -76,10 +77,16 @@ jobs:
7677
echo "Error: requirements_syncKnowledgeBaseFunction is empty or contains only blank lines"; \
7778
exit 1; \
7879
fi
80+
if [ ! -s requirements_bedrockLoggingConfigFunction ] || [ "$(grep -c -v '^[[:space:]]*$' requirements_bedrockLoggingConfigFunction)" -eq 0 ]; then \
81+
echo "Error: requirements_bedrockLoggingConfigFunction is empty or contains only blank lines"; \
82+
exit 1; \
83+
fi
7984
mkdir -p .dependencies/slackBotFunction/python
8085
mkdir -p .dependencies/syncKnowledgeBaseFunction/python
86+
mkdir -p .dependencies/bedrockLoggingConfigFunction/python
8187
pip3 install -r requirements_slackBotFunction -t .dependencies/slackBotFunction/python
8288
pip3 install -r requirements_syncKnowledgeBaseFunction -t .dependencies/syncKnowledgeBaseFunction/python
89+
pip3 install -r requirements_bedrockLoggingConfigFunction -t .dependencies/bedrockLoggingConfigFunction/python
8390
8491
- name: "Tar files"
8592
run: |

‎Makefile‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ lint-flake8:
4848
test:
4949
cd packages/slackBotFunction && PYTHONPATH=. COVERAGE_FILE=coverage/.coverage poetry run python -m pytest
5050
cd packages/syncKnowledgeBaseFunction && PYTHONPATH=. COVERAGE_FILE=coverage/.coverage poetry run python -m pytest
51+
cd packages/bedrockLoggingConfigFunction && PYTHONPATH=. COVERAGE_FILE=coverage/.coverage poetry run python -m pytest
5152

5253
clean:
5354
rm -rf packages/cdk/coverage
@@ -104,6 +105,7 @@ cdk-deploy: guard-STACK_NAME
104105
cdk-synth:
105106
mkdir -p .dependencies/slackBotFunction
106107
mkdir -p .dependencies/syncKnowledgeBaseFunction
108+
mkdir -p .dependencies/bedrockLoggingConfigFunction
107109
mkdir -p .local_config
108110
STACK_NAME=epsam \
109111
COMMIT_ID=undefined \
Lines changed: 186 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,186 @@
1+
"""
2+
cloudformation has no native resource for bedrock model invocation logging
3+
this custom resource bridges that gap via the bedrock api
4+
"""
5+
6+
import json
7+
import os
8+
import traceback
9+
import boto3
10+
import urllib3
11+
from aws_lambda_powertools import Logger
12+
13+
http = urllib3.PoolManager()
14+
logger = Logger()
15+
16+
17+
def send_response(event, context, response_status, response_data, physical_resource_id=None, reason=None):
18+
"""
19+
signals cloudformation that the custom resource operation completed
20+
"""
21+
response_url = event["ResponseURL"]
22+
23+
response_body = {
24+
"Status": response_status,
25+
"Reason": reason or f"See CloudWatch Log Stream: {context.log_stream_name}",
26+
"PhysicalResourceId": physical_resource_id or context.log_stream_name,
27+
"StackId": event["StackId"],
28+
"RequestId": event["RequestId"],
29+
"LogicalResourceId": event["LogicalResourceId"],
30+
"Data": response_data,
31+
}
32+
33+
json_response_body = json.dumps(response_body)
34+
35+
headers = {"content-type": "", "content-length": str(len(json_response_body))}
36+
37+
try:
38+
http.request("PUT", response_url, body=json_response_body, headers=headers)
39+
logger.info(f"cloudformation response sent: {response_status}")
40+
except Exception as e:
41+
logger.error(f"failed to signal cloudformation: {str(e)}")
42+
43+
44+
def parse_event(event):
45+
"""parse event to determine request type and logging state"""
46+
is_direct_invocation = not event or "RequestType" not in event
47+
48+
if is_direct_invocation:
49+
logger.info("direct invocation detected - treating as Update operation")
50+
request_type = "Update"
51+
52+
# check for enable_logging override in event, otherwise use env var
53+
if "enable_logging" in event:
54+
enable_logging = str(event.get("enable_logging", "true")).lower() == "true"
55+
logger.info(f"using enable_logging from event payload: {enable_logging}")
56+
else:
57+
enable_logging = os.environ.get("ENABLE_LOGGING", "true").lower() == "true"
58+
logger.info(f"using ENABLE_LOGGING from environment: {enable_logging}")
59+
else:
60+
# cloudformation invocation - always use env var
61+
request_type = event["RequestType"]
62+
enable_logging = os.environ.get("ENABLE_LOGGING", "true").lower() == "true"
63+
logger.info(f"cloudformation invocation - using ENABLE_LOGGING from environment: {enable_logging}")
64+
65+
return request_type, enable_logging, is_direct_invocation
66+
67+
68+
def handle_logging_disabled(event, context, bedrock, is_direct_invocation):
69+
"""handle case when logging is disabled"""
70+
logger.info("bedrock logging disabled - removing configuration")
71+
try:
72+
bedrock.delete_model_invocation_logging_configuration()
73+
logger.info("bedrock logging configuration deleted")
74+
except bedrock.exceptions.ResourceNotFoundException:
75+
logger.info("logging configuration not found (already disabled)")
76+
77+
# only send cloudformation response if this is a real cfn event
78+
if not is_direct_invocation:
79+
send_response(
80+
event,
81+
context,
82+
"SUCCESS",
83+
{"Message": "Bedrock logging disabled via environment variable"},
84+
physical_resource_id="BedrockModelInvocationLogging",
85+
)
86+
87+
88+
def handle_create_or_update(event, context, bedrock, is_direct_invocation):
89+
"""handle create or update operations"""
90+
logger.info("configuring bedrock model invocation logging")
91+
92+
# Get CloudWatch config from environment variables (set by CDK)
93+
cloudwatch_log_group_name = os.environ.get("CLOUDWATCH_LOG_GROUP_NAME")
94+
cloudwatch_role_arn = os.environ.get("CLOUDWATCH_ROLE_ARN")
95+
96+
# aws requires at least one logging destination
97+
if not cloudwatch_log_group_name or not cloudwatch_role_arn:
98+
error_msg = """
99+
CLOUDWATCH_LOG_GROUP_NAME and CLOUDWATCH_ROLE_ARN environment variables required.
100+
Cannot configure logging without destination."""
101+
102+
logger.error(error_msg)
103+
if is_direct_invocation:
104+
raise ValueError(error_msg)
105+
send_response(event, context, "FAILED", {}, reason=error_msg)
106+
return
107+
108+
logging_config = {
109+
"cloudWatchConfig": {
110+
"logGroupName": cloudwatch_log_group_name,
111+
"roleArn": cloudwatch_role_arn,
112+
},
113+
}
114+
115+
logger.info(f"cloudwatch logs enabled: {cloudwatch_log_group_name}")
116+
117+
response = bedrock.put_model_invocation_logging_configuration(loggingConfig=logging_config)
118+
logger.info(f"bedrock logging configured: {json.dumps(response)}")
119+
120+
# only send cloudformation response if this is a real cfn event
121+
if not is_direct_invocation:
122+
send_response(
123+
event,
124+
context,
125+
"SUCCESS",
126+
{
127+
"Message": "Bedrock model invocation logging configured successfully",
128+
"CloudWatchLogGroup": cloudwatch_log_group_name,
129+
},
130+
physical_resource_id="BedrockModelInvocationLogging",
131+
)
132+
133+
134+
def handle_delete(event, context, bedrock, is_direct_invocation):
135+
"""handle delete operations"""
136+
logger.info("deleting bedrock model invocation logging")
137+
138+
try:
139+
bedrock.delete_model_invocation_logging_configuration()
140+
logger.info("bedrock logging configuration deleted")
141+
except bedrock.exceptions.ResourceNotFoundException:
142+
logger.info("logging configuration not found")
143+
144+
# only send cloudformation response if this is a real cfn event
145+
if not is_direct_invocation:
146+
send_response(
147+
event,
148+
context,
149+
"SUCCESS",
150+
{"Message": "Bedrock model invocation logging deleted successfully"},
151+
physical_resource_id="BedrockModelInvocationLogging",
152+
)
153+
154+
155+
@logger.inject_lambda_context(log_event=True, clear_state=True)
156+
def handler(event, context):
157+
"""
158+
configures bedrock model invocation logging via put/delete api calls
159+
toggleable via ENABLE_LOGGING environment variable
160+
161+
supports direct invocation in aws console with:
162+
- {} (empty) - uses ENABLE_LOGGING env var
163+
- {"enable_logging": true/false} - overrides env var
164+
"""
165+
request_type, enable_logging, is_direct_invocation = parse_event(event)
166+
167+
bedrock = boto3.client("bedrock")
168+
169+
try:
170+
if request_type in ["Create", "Update"]:
171+
if not enable_logging:
172+
handle_logging_disabled(event, context, bedrock, is_direct_invocation)
173+
return
174+
handle_create_or_update(event, context, bedrock, is_direct_invocation)
175+
elif request_type == "Delete":
176+
handle_delete(event, context, bedrock, is_direct_invocation)
177+
else:
178+
if not is_direct_invocation:
179+
send_response(event, context, "FAILED", {}, reason=f"unsupported request type: {request_type}")
180+
except Exception as e:
181+
error_message = f"error: {str(e)}\n{traceback.format_exc()}"
182+
logger.error(error_message)
183+
if not is_direct_invocation:
184+
send_response(event, context, "FAILED", {}, reason=error_message)
185+
else:
186+
raise # re-raise for direct invocations so user sees the error
Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
[pytest]
2+
testpaths = tests
3+
python_files = test_*.py
4+
python_functions = test_*
5+
addopts = -v --tb=short --cov=app --cov-report=xml:coverage/coverage.xml --cov-report=term-missing --cov-config=pytest.ini
6+
7+
[coverage:run]
8+
omit = */__init__.py

‎packages/bedrockLoggingConfigFunction/tests/__init__.py‎

Whitespace-only changes.

0 commit comments

Comments
 (0)