|
4 | 4 | import inspect |
5 | 5 | import io |
6 | 6 | import json |
| 7 | +import logging |
7 | 8 | import math |
8 | 9 | from typing import Any, Dict, Optional |
9 | 10 |
|
10 | 11 | from botocore.response import StreamingBody |
11 | 12 |
|
| 13 | +_logger = logging.getLogger(__name__) |
| 14 | +# Set logger level to DEBUG |
| 15 | +_logger.setLevel(logging.DEBUG) |
| 16 | + |
12 | 17 | from amazon.opentelemetry.distro._aws_attribute_keys import ( |
13 | 18 | AWS_BEDROCK_AGENT_ID, |
14 | 19 | AWS_BEDROCK_DATA_SOURCE_ID, |
@@ -275,7 +280,7 @@ def extract_attributes(self, attributes: _AttributeMapT): |
275 | 280 | self._extract_mistral_attributes(attributes, request_body) |
276 | 281 |
|
277 | 282 | except json.JSONDecodeError: |
278 | | - print("Error: Unable to parse the body as JSON") |
| 283 | + _logger.debug("Error: Unable to parse the body as JSON") |
279 | 284 |
|
280 | 285 | def _extract_titan_attributes(self, attributes, request_body): |
281 | 286 | config = request_body.get("textGenerationConfig", {}) |
@@ -352,9 +357,9 @@ def on_success(self, span: Span, result: Dict[str, Any]): |
352 | 357 | result["body"] = StreamingBody(new_stream, len(body_content)) |
353 | 358 |
|
354 | 359 | except json.JSONDecodeError: |
355 | | - print("Error: Unable to parse the response body as JSON") |
| 360 | + _logger.debug("Error: Unable to parse the response body as JSON") |
356 | 361 | except Exception as e: # pylint: disable=broad-exception-caught, invalid-name |
357 | | - print(f"Error processing response: {str(e)}") |
| 362 | + _logger.debug(f"Error processing response: {str(e)}") |
358 | 363 | finally: |
359 | 364 | if original_body is not None: |
360 | 365 | original_body.close() |
|
0 commit comments