feat(adk): add native Bedrock embedding support for agent memory#1641
Merged
EItanya merged 4 commits intokagent-dev:mainfrom Apr 10, 2026
Merged
feat(adk): add native Bedrock embedding support for agent memory#1641EItanya merged 4 commits intokagent-dev:mainfrom
EItanya merged 4 commits intokagent-dev:mainfrom
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds first-class AWS Bedrock (Titan) embedding support to KagentMemoryService, allowing agent memory embeddings to work when EmbeddingConfig.provider="bedrock" without falling back to an OpenAI-compatible client.
Changes:
- Add a Bedrock embedding implementation (
_embed_bedrock) and wire it into the embedding-provider dispatcher. - Use
boto3.client("bedrock-runtime").invoke_model()per text (parallelized viaasyncio.gather+to_thread). - Add unit tests covering Bedrock embedding, region selection via env vars, and error handling.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
python/packages/kagent-adk/src/kagent/adk/_memory_service.py |
Adds Bedrock embedding provider dispatch + _embed_bedrock() implementation using boto3 Bedrock Runtime. |
python/packages/kagent-adk/tests/unittests/test_embedding.py |
Adds unit tests validating Bedrock embedding behavior, env-region handling, and graceful failure. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Add _embed_bedrock() to KagentMemoryService so that memory works with AWS Bedrock embedding models (e.g. amazon.titan-embed-text-v1/v2) without falling back to the OpenAI-compatible client. The implementation uses boto3's invoke_model API, consistent with how KAgentBedrockLlm already handles LLM calls. Region is resolved from AWS_DEFAULT_REGION / AWS_REGION env vars (same credential chain as the Bedrock LLM provider). Each text is embedded individually since the Titan Embedding API accepts a single inputText per invocation; asyncio.gather parallelises the calls. The upstream truncation-to-768 + L2-normalization logic handles dimension differences across Titan v1 (1536d) and v2 (1024d) models, so no model-specific dimension parameter is needed. Signed-off-by: Gonçalo Faustino <goncalo.santos@wellhub.com> Made-with: Cursor
- Remove unused mock_response in test_bedrock_embed - Override AWS_DEFAULT_REGION in region env test to prevent flakiness when the runner environment has it set Signed-off-by: Gonçalo Faustino <goncalo.santos@wellhub.com> Made-with: Cursor
Wrap long mock_client.invoke_model lines to comply with ruff line-length formatting rules. Signed-off-by: Gonçalo Faustino <goncalo.santos@wellhub.com> Made-with: Cursor
da00b95 to
6bd89de
Compare
EItanya
approved these changes
Apr 10, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds native AWS Bedrock embedding support to
KagentMemoryService, so agent memory works out of the box with Bedrock embedding models (e.g.amazon.titan-embed-text-v1,amazon.titan-embed-text-v2:0).Currently, when
provider="bedrock"is set on an embeddingModelConfig, the memory service falls back to an OpenAI-compatible client which fails because noOPENAI_API_KEYis available:Changes
_memory_service.py: Add_embed_bedrock()method and register"bedrock"in_call_embedding_providerdispatchtest_embedding.py: Add 3 unit tests (basic embedding, region from env, error handling)Implementation details
boto3.client("bedrock-runtime").invoke_model(), consistent with howKAgentBedrockLlmhandles LLM calls inmodels/_bedrock.pyAWS_DEFAULT_REGION/AWS_REGIONenv vars (same credential chain as the Bedrock LLM provider — IRSA, instance profile, etc.)inputTextper invocation;asyncio.gatherparallelises the calls_generate_embedding_asynchandles dimension differences across Titan v1 (1536d) and v2 (1024d) models, so no model-specificdimensionsparameter is neededboto3is already a declared dependency ofkagent-adkTesting
New tests:
test_bedrock_embed— verifies embedding generation and truncation to 768 dimensionstest_bedrock_embed_uses_region_from_env— verifiesAWS_REGIONenv var is passed to boto3 clienttest_bedrock_embed_error_returns_empty— verifies graceful error handling returns[]