diff --git a/.changelog/4997.fixed b/.changelog/4997.fixed new file mode 100644 index 0000000000..b1228ff155 --- /dev/null +++ b/.changelog/4997.fixed @@ -0,0 +1 @@ +`opentelemetry-instrumentation-botocore`: tag Bedrock embedding InvokeModel calls as embeddings diff --git a/instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock.py b/instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock.py index a4f8f05fad..eaf4d0fb92 100644 --- a/instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock.py +++ b/instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock.py @@ -100,6 +100,12 @@ _MODEL_ID_KEY: str = "modelId" +def _is_embedding_model(model_id: str) -> bool: + if not model_id: + return False + return "embed" in model_id.rsplit(".", 1)[-1].lower() + + class _BedrockRuntimeExtension(_AwsSdkExtension): """ This class is an extension for _AttributeMapT: attributes[GEN_AI_REQUEST_MODEL] = model_id - # titan in invoke model is a text completion one - if "body" in self._call_context.params and "amazon.titan" in model_id: + if _is_embedding_model(model_id): + attributes[GEN_AI_OPERATION_NAME] = GenAiOperationNameValues.EMBEDDINGS.value + elif "body" in self._call_context.params and "amazon.titan" in model_id: + # titan in invoke model is a text completion one attributes[GEN_AI_OPERATION_NAME] = GenAiOperationNameValues.TEXT_COMPLETION.value else: attributes[GEN_AI_OPERATION_NAME] = GenAiOperationNameValues.CHAT.value @@ -164,7 +172,10 @@ def extract_attributes(self, attributes: _AttributeMapT): model_id = self._call_context.params.get(_MODEL_ID_KEY) if model_id: attributes[GEN_AI_REQUEST_MODEL] = model_id - attributes[GEN_AI_OPERATION_NAME] = GenAiOperationNameValues.CHAT.value + if _is_embedding_model(model_id): + attributes[GEN_AI_OPERATION_NAME] = GenAiOperationNameValues.EMBEDDINGS.value + else: + attributes[GEN_AI_OPERATION_NAME] = GenAiOperationNameValues.CHAT.value # Converse / ConverseStream if inference_config := self._call_context.params.get("inferenceConfig"): @@ -196,7 +207,9 @@ def extract_attributes(self, attributes: _AttributeMapT): try: request_body = json.loads(body) - if "amazon.titan" in model_id: + if _is_embedding_model(model_id): + attributes[GEN_AI_OPERATION_NAME] = GenAiOperationNameValues.EMBEDDINGS.value + elif "amazon.titan" in model_id: # titan interface is a text completion one attributes[GEN_AI_OPERATION_NAME] = GenAiOperationNameValues.TEXT_COMPLETION.value self._extract_titan_attributes(attributes, request_body) @@ -357,7 +370,7 @@ def _get_request_messages(self): # if no messages interface, convert to messages format from generic API if not messages: model_id = self._call_context.params.get(_MODEL_ID_KEY) - if "amazon.titan" in model_id: + if "amazon.titan" in model_id and not _is_embedding_model(model_id): messages = self._get_messages_from_input_text(decoded_body, "inputText") elif "cohere.command-r" in model_id: # chat_history can be converted to messages; for now, just use message @@ -474,6 +487,8 @@ def _invoke_model_on_success( result["body"] = StreamingBody(new_stream, len(body_content)) response_body = json.loads(body_content.decode("utf-8")) + if _is_embedding_model(model_id): + return if "amazon.titan" in model_id: self._handle_amazon_titan_response(span, response_body, instrumentor_context, capture_content) elif "amazon.nova" in model_id: diff --git a/instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_embeddings_model[amazon.titan].yaml b/instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_embeddings_model[amazon.titan].yaml new file mode 100644 index 0000000000..6f4cc099d9 --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_embeddings_model[amazon.titan].yaml @@ -0,0 +1,58 @@ +interactions: +- request: + body: |- + { + "inputText": "Say this is a test" + } + headers: + Content-Length: + - '36' + User-Agent: + - Boto3/1.35.56 md/Botocore#1.35.56 ua/2.0 os/macos#24.0.0 md/arch#arm64 lang/python#3.13.1 + md/pyimpl#CPython cfg/retry-mode#legacy Botocore/1.35.56 + X-Amz-Date: + - 20250826T103000Z + X-Amzn-Trace-Id: + - Root=1-aaaaaaaa-bbbbbbbbbbbbbbbbbbbbbbbb;Parent=cccccccccccccccc;Sampled=1 + amz-sdk-invocation-id: + - 11111111-2222-3333-4444-555555555555 + amz-sdk-request: + - attempt=1 + authorization: + - Bearer test_aws_authorization + method: POST + uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-embed-text-v1/invoke + response: + body: + string: |- + { + "embedding": [ + 0.0135, + -0.0272, + 0.0047, + 0.0189, + -0.0091, + 0.0023, + 0.0314, + -0.0168 + ], + "inputTextTokenCount": 5 + } + headers: + Connection: + - keep-alive + Content-Type: + - application/json + Date: + - Wed, 26 Aug 2026 10:30:00 GMT + Set-Cookie: test_set_cookie + X-Amzn-Bedrock-Input-Token-Count: + - '5' + X-Amzn-Bedrock-Invocation-Latency: + - '120' + x-amzn-RequestId: + - aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee + status: + code: 200 + message: OK +version: 1 diff --git a/instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_embeddings_model[cohere.embed].yaml b/instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_embeddings_model[cohere.embed].yaml new file mode 100644 index 0000000000..7bed13865f --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_embeddings_model[cohere.embed].yaml @@ -0,0 +1,67 @@ +interactions: +- request: + body: |- + { + "texts": [ + "Say this is a test" + ], + "input_type": "search_document" + } + headers: + Content-Length: + - '76' + User-Agent: + - Boto3/1.35.56 md/Botocore#1.35.56 ua/2.0 os/macos#24.0.0 md/arch#arm64 lang/python#3.13.1 + md/pyimpl#CPython cfg/retry-mode#legacy Botocore/1.35.56 + X-Amz-Date: + - 20250826T103001Z + X-Amzn-Trace-Id: + - Root=1-ffffffff-eeeeeeeeeeeeeeeeeeeeeeee;Parent=dddddddddddddddd;Sampled=1 + amz-sdk-invocation-id: + - 66666666-7777-8888-9999-000000000000 + amz-sdk-request: + - attempt=1 + authorization: + - Bearer test_aws_authorization + method: POST + uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/cohere.embed-v4%3A0/invoke + response: + body: + string: |- + { + "id": "836a33cc-61ec-4e65-afaf-c4628171a315", + "embeddings": [ + [ + 0.0135, + -0.0272, + 0.0047, + 0.0189, + -0.0091, + 0.0023, + 0.0314, + -0.0168 + ] + ], + "response_type": "embeddings_floats", + "texts": [ + "Say this is a test" + ] + } + headers: + Connection: + - keep-alive + Content-Type: + - application/json + Date: + - Wed, 26 Aug 2026 10:30:01 GMT + Set-Cookie: test_set_cookie + X-Amzn-Bedrock-Input-Token-Count: + - '5' + X-Amzn-Bedrock-Invocation-Latency: + - '95' + x-amzn-RequestId: + - 836a33cc-61ec-4e65-afaf-c4628171a315 + status: + code: 200 + message: OK +version: 1 diff --git a/instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_bedrock.py b/instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_bedrock.py index f90144a851..a9b0750d10 100644 --- a/instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_bedrock.py +++ b/instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_bedrock.py @@ -1269,6 +1269,68 @@ def get_model_name_from_family(llm_model): return llm_model_name[llm_model] +def get_embeddings_model_name_from_family(model_family): + return { + "amazon.titan": "amazon.titan-embed-text-v1", + "cohere.embed": "cohere.embed-v4:0", + }[model_family] + + +def get_invoke_embeddings_body(llm_model): + if "cohere.embed" in llm_model: + return json.dumps( + { + "texts": ["Say this is a test"], + "input_type": "search_document", + } + ) + if "amazon.titan" in llm_model: + return json.dumps({"inputText": "Say this is a test"}) + raise ValueError(f"No embeddings config for {llm_model}") + + +@pytest.mark.parametrize( + "model_family", + [ + "amazon.titan", + "cohere.embed", + ], +) +@pytest.mark.vcr() +def test_invoke_model_with_embeddings_model( + span_exporter, + log_exporter, + bedrock_runtime_client, + instrument_with_content, + model_family, +): + llm_model_value = get_embeddings_model_name_from_family(model_family) + body = get_invoke_embeddings_body(llm_model_value) + response = bedrock_runtime_client.invoke_model( + body=body, + modelId=llm_model_value, + ) + + response_body = json.loads(response["body"].read()) + assert response_body + if model_family == "amazon.titan": + assert "embedding" in response_body + elif model_family == "cohere.embed": + assert "embeddings" in response_body + else: + pytest.xfail(f"model family not handled: {model_family}") + + (span,) = span_exporter.get_finished_spans() + assert_stream_completion_attributes( + span, + llm_model_value, + operation_name="embeddings", + ) + + logs = log_exporter.get_finished_logs() + assert len(logs) == 0 + + @pytest.mark.parametrize( "model_family", [