Map top_p to topP in the Bedrock Titan and AI21 completions adapters - #25795

Related Issues/PRs

Relates to #25571

What changes are proposed in this pull request?

Both AWSTitanAdapter.completions_to_model and AI21Adapter.completions_to_model build their Bedrock request body with rename_payload_keys, which passes through any key absent from its mapping unchanged. Neither mapping covered top_p, so a caller-supplied value reached Bedrock under the literal key top_p. Titan's textGenerationConfig and Jurassic's request body both name that field topP, so Bedrock did not recognise the key and silently ignored it — the request succeeded and returned a normal completion, just without the nucleus sampling the caller asked for.

This adds "top_p": "topP" to both mappings. The outbound body, produced by running the adapters:

On top_k — flagging a deviation from the triage note so it is cheap to redirect

The triage note on #25571 suggested mapping top_k -> topK alongside top_p. I did not, because I could not find a topK field on either model family to map onto:

• Titan's textGenerationConfig accepts only maxTokenCount, stopSequences, temperature and topP.
• Jurassic exposes topKReturn, which controls how many alternative tokens are reported, not top-k sampling. Mapping top_k onto it would change the response payload rather than the sampling behaviour.
• Corroborating from inside this repo: the Converse path in this same file builds inferenceConfig from temperature/topP/maxTokens/stopSequences and omits topK (bedrock.py:408-421), and no Bedrock code path in the repo maps topK anywhere. GeminiAdapter does map top_k -> topK, because Gemini genuinely supports it.

Renaming top_k -> topK here would swap one key Bedrock ignores for another while making the code look like it handles the parameter — the same silent-drop failure mode the issue is about. So instead of deleting the two # TODO handle top_p, top_k, etc. comments as though they were fully addressed, I replaced them with notes recording why top_k and the penalty parameters stay unmapped.

That leaves a decision I would rather you make than guess at. top_k is currently accepted and ignored; the alternative is to reject it with a 422, which this file already does for n != 1 on Titan (and AnthropicAdapter does for n, and for temperature+top_p together). I left that out because it turns a currently-succeeding request into a hard error and CONTRIBUTING.md asks for backwards-compatible changes — but it is a small follow-up in either direction, and I am happy to add it here if you prefer. I have kept this as "Relates to" rather than "Closes" so #25571 does not auto-close while that is open.

How is this PR tested?

• Existing unit/integration tests
• New unit/integration tests

Two new cases in bedrock_model_provider_fixtures (Titan and Jurassic, each sending top_p). These run through the existing test_bedrock_request_response, which drives the real AmazonBedrockProvider.completions and asserts the exact body handed to Bedrock via mock_request.assert_called_once_with(model_request) — so they pin the outbound key name, not just the adapter's return value. Each fixture runs against all four bedrock_aws_configs.

I also confirmed the new tests actually catch the bug, by reverting only the two mapping entries and re-running. They fail with exactly the right message, and the six pre-existing cases still pass:

AssertionError: expected call not found.
Expected: _request({'inputText': 'This is a test', 'textGenerationConfig': {'maxTokenCount': 1000, 'topP': 0.9}})
  Actual: _request({'inputText': 'This is a test', 'textGenerationConfig': {'maxTokenCount': 1000, 'top_p': 0.9}})

[3] AMAZON     amazon.titan-tg1-large   -> PASS
[4] AMAZON     amazon.titan-tg1-large   -> PASS  <-- top_p case
      outbound body: {'inputText': 'This is a test', 'textGenerationConfig': {'maxTokenCount': 1000, 'topP': 0.9}}
[5] AI21       ai21.j2-ultra            -> PASS
[6] AI21       ai21.j2-ultra            -> PASS  <-- top_p case
      outbound body: {'prompt': 'This is a test', 'maxTokens': 1000, 'topP': 0.9}
[7] AI21       ai21.j2-mid              -> PASS

total=32 passed=32 skipped=0

ruff check and ruff format --check pass on both changed files. Commits are DCO signed off.

Does this PR require documentation update?

• No.

Does this PR require updating the MLflow Skills repository?

• No.

Release Notes

Is this a user-facing change?

• Yes. Give a description of this change to be included in the release notes for MLflow users.

Fixed top_p being silently ignored for Amazon Titan and AI21 Jurassic models on the AI Gateway's llm/v1/completions endpoint. The parameter was forwarded to Bedrock under its OpenAI-style top_p name instead of the topP name those models expect, so requests succeeded while ignoring the caller's nucleus-sampling setting.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

• area/gateway: MLflow AI Gateway client APIs, server, and third-party integrations

How should the PR be classified in the release notes? Choose one:

• rn/bug-fix - A user-facing bug fix worth mentioning in the release notes

Is this PR a critical bugfix or security fix that should go into the next patch release?

• This PR is critical and needs to be in the next patch release
• This PR can wait for the next minor release

Disclosure: I used an AI assistant (Claude) to help trace the call paths, implement the fix, and draft this description. Every file/line reference and every result above was produced by running the code, and I reviewed the change before submitting.

🤖 Generated with Claude Code