Skip to content

Commit

Permalink
Merge branch 'main' into add-checksum
Browse files Browse the repository at this point in the history
  • Loading branch information
harrryr authored Dec 17, 2024
2 parents b3fe88e + 93c1588 commit bfdfee9
Show file tree
Hide file tree
Showing 5 changed files with 304 additions and 196 deletions.
2 changes: 1 addition & 1 deletion Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ RUN sed -i "/opentelemetry-exporter-otlp-proto-grpc/d" ./aws-opentelemetry-distr
RUN mkdir workspace && pip install --target workspace ./aws-opentelemetry-distro

# Stage 2: Build the cp-utility binary
FROM public.ecr.aws/docker/library/rust:1.75 as builder
FROM public.ecr.aws/docker/library/rust:1.81 as builder

WORKDIR /usr/src/cp-utility
COPY ./tools/cp-utility .
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@
_GUARDRAIL_ID: str = "guardrailId"
_GUARDRAIL_ARN: str = "guardrailArn"
_MODEL_ID: str = "modelId"
_AWS_BEDROCK_SYSTEM: str = "aws_bedrock"
_AWS_BEDROCK_SYSTEM: str = "aws.bedrock"

_logger = logging.getLogger(__name__)
# Set logger level to DEBUG
Expand Down Expand Up @@ -268,6 +268,8 @@ def extract_attributes(self, attributes: _AttributeMapT):

if "amazon.titan" in model_id:
self._extract_titan_attributes(attributes, request_body)
if "amazon.nova" in model_id:
self._extract_nova_attributes(attributes, request_body)
elif "anthropic.claude" in model_id:
self._extract_claude_attributes(attributes, request_body)
elif "meta.llama" in model_id:
Expand All @@ -288,6 +290,12 @@ def _extract_titan_attributes(self, attributes, request_body):
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, config.get("topP"))
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, config.get("maxTokenCount"))

def _extract_nova_attributes(self, attributes, request_body):
config = request_body.get("inferenceConfig", {})
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, config.get("temperature"))
self._set_if_not_none(attributes, GEN_AI_REQUEST_TOP_P, config.get("top_p"))
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, config.get("max_new_tokens"))

def _extract_claude_attributes(self, attributes, request_body):
self._set_if_not_none(attributes, GEN_AI_REQUEST_MAX_TOKENS, request_body.get("max_tokens"))
self._set_if_not_none(attributes, GEN_AI_REQUEST_TEMPERATURE, request_body.get("temperature"))
Expand Down Expand Up @@ -324,6 +332,7 @@ def _set_if_not_none(attributes, key, value):
if value is not None:
attributes[key] = value

# pylint: disable=too-many-branches
def on_success(self, span: Span, result: Dict[str, Any]):
model_id = self._call_context.params.get(_MODEL_ID)

Expand All @@ -342,6 +351,8 @@ def on_success(self, span: Span, result: Dict[str, Any]):
response_body = json.loads(telemetry_content.decode("utf-8"))
if "amazon.titan" in model_id:
self._handle_amazon_titan_response(span, response_body)
if "amazon.nova" in model_id:
self._handle_amazon_nova_response(span, response_body)
elif "anthropic.claude" in model_id:
self._handle_anthropic_claude_response(span, response_body)
elif "meta.llama" in model_id:
Expand Down Expand Up @@ -375,6 +386,17 @@ def _handle_amazon_titan_response(self, span: Span, response_body: Dict[str, Any
if "completionReason" in result:
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [result["completionReason"]])

# pylint: disable=no-self-use
def _handle_amazon_nova_response(self, span: Span, response_body: Dict[str, Any]):
if "usage" in response_body:
usage = response_body["usage"]
if "inputTokens" in usage:
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, usage["inputTokens"])
if "outputTokens" in usage:
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, usage["outputTokens"])
if "stopReason" in response_body:
span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [response_body["stopReason"]])

# pylint: disable=no-self-use
def _handle_anthropic_claude_response(self, span: Span, response_body: Dict[str, Any]):
if "usage" in response_body:
Expand Down
Loading

0 comments on commit bfdfee9

Please sign in to comment.