Skip to content

Commit

Permalink
Make it explicit we are handling only the converse operation
Browse files Browse the repository at this point in the history
  • Loading branch information
xrmx committed Dec 31, 2024
1 parent c351b0d commit fecd5d4
Showing 1 changed file with 11 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -53,9 +53,14 @@ class _BedrockRuntimeExtension(_AwsSdkExtension):
Amazon Bedrock Runtime</a>.
"""

_HANDLED_OPERATIONS = {"Converse"}

def extract_attributes(self, attributes: _AttributeMapT):
attributes[GEN_AI_SYSTEM] = GenAiSystemValues.AWS_BEDROCK.value

if self._call_context.operation not in self._HANDLED_OPERATIONS:
return

model_id = self._call_context.params.get(_MODEL_ID_KEY)
if model_id:
attributes[GEN_AI_REQUEST_MODEL] = model_id
Expand Down Expand Up @@ -97,6 +102,9 @@ def _set_if_not_none(attributes, key, value):
attributes[key] = value

def before_service_call(self, span: Span):
if self._call_context.operation not in self._HANDLED_OPERATIONS:
return

if not span.is_recording():
return

Expand All @@ -107,6 +115,9 @@ def before_service_call(self, span: Span):
span.update_name(f"{operation_name} {request_model}")

def on_success(self, span: Span, result: dict[str, Any]):
if self._call_context.operation not in self._HANDLED_OPERATIONS:
return

model_id = self._call_context.params.get(_MODEL_ID_KEY)

if not model_id:
Expand Down

0 comments on commit fecd5d4

Please sign in to comment.