fix(botocore): tag Bedrock embedding InvokeModel calls as embeddings - #4997
fix(botocore): tag Bedrock embedding InvokeModel calls as embeddings#4997BetterAndBetterII wants to merge 4 commits into
Conversation
Pull request dashboard statusWaiting on the author · refreshed 2026-08-29 10:21 UTC Respond to 4 review items (e.g. link a commit, explain why not, ask a follow-up): Status above doesn't look right?
|
|
Not an approver, but I checked this out and ran it on Python 3.13. test_botocore_bedrock.py is 91/91 green on the branch. Reverting only bedrock.py to main while keeping the new tests fails all 16 embedding cases, so the tests cover the behavior they claim. GenAiOperationNameValues.EMBEDDINGS.value is "embeddings" in the semconv package this code imports, matching the spec's embeddings span definition. The pre-fix failures show the bug had two shapes. Titan embed models returned chat with no body and text_completion with a body, so the same model produced different operation names depending on call shape. Worth adding to the description.
|
|
|
||
| @pytest.mark.parametrize("model_id", EMBEDDING_MODEL_IDS) | ||
| @pytest.mark.parametrize("with_body", [False, True]) | ||
| def test_extract_attributes_embedding_operation_name(model_id, with_body): |
There was a problem hiding this comment.
I would prefer to have recorded tests instead of unit testing the helpers. I don't think we have any embedding test already, so introducing something like a parametrized test_invoke_model_with_embeddings_model. Also I think one model entry per model family is enough.
There was a problem hiding this comment.
Replaced the helper unit tests with parametrized recorded test_invoke_model_with_embeddings_model (one model per family).
|
|
||
|
|
||
| def _is_embedding_model(model_id: str) -> bool: | ||
| return "embed" in model_id.lower() |
There was a problem hiding this comment.
Are model names case insensitive? Anyway it looks like this will work from what I can see in https://docs.aws.amazon.com/bedrock/latest/userguide/models-api-compatibility.html
There was a problem hiding this comment.
As @JoeyPatricio spotted maybe we should split the model_id by . and look only in the last part?
There was a problem hiding this comment.
Detection now uses only the last dotted segment of modelId.
There was a problem hiding this comment.
AWS docs list model IDs in lowercase; the last-segment check still lowercases so mixed-case IDs match.
github-project-automation
Bot
moved this to Reviewed PRs that need fixes
in Python PR digest
Detect embeddings from the last modelId segment and add recorded InvokeModel tests for Titan and Cohere embed families.
62c025c to
e6f20be
Compare
| result["body"] = StreamingBody(new_stream, len(body_content)) | ||
|
|
||
| response_body = json.loads(body_content.decode("utf-8")) | ||
| if _is_embedding_model(model_id): |
There was a problem hiding this comment.
IIUC this makes gen_ai.usage.input_tokens disappear for Titan embed models on main the response reaches _handle_amazon_titan_response, which sets it from inputTextTokenCount.
| else: | ||
| pytest.xfail(f"model family not handled: {model_family}") | ||
|
|
||
| (span,) = span_exporter.get_finished_spans() |
There was a problem hiding this comment.
The new test also pins this: assert_stream_completion_attributes leaves input_tokens as None
henry3260
left a comment
henry3260
left a comment
There was a problem hiding this comment.
Non-blocking thought: this fix had to touch 5 places because the model_id dispatch is repeated across 4 functions (_extract_metrics_attributes, extract_attributes, _get_request_messages, _invoke_model_on_success). Would it be worth centralizing the family detection?
4 participants
Description
extract_attributesand_extract_metrics_attributesboth labeled every BedrockInvokeModelcall aschat(or Titantext_completionwhen themodelIdcontainedamazon.titan). Embedding model IDs such ascohere.embed-v4:0andamazon.titan-embed-*were therefore recorded withgen_ai.operation.name=chat, and the span name plus duration/token metrics inherited that value.Detect embedding
modelIds (substringembed) first in both extract paths and setgen_ai.operation.nametoembeddings. Titan text-completion and chat models are unchanged.Fixes #4996
Type of change
How Has This Been Tested?
test_botocore_bedrock.pycovering both extract sites, with and without a request body, for Cohere and Titan embedding model IDs, plus unchanged Titan text-completion and chat modelsDoes This PR Require a Core Repo Change?
Checklist: