feat: add Managed Knowledge Base support to AWS Bedrock KB plugin#3425
feat: add Managed Knowledge Base support to AWS Bedrock KB plugin#3425PVidyadhar wants to merge 1 commit into
Conversation
- Added Knowledge Base Type dropdown (MANAGED/VECTOR) to provider config - MANAGED path uses managedSearchConfiguration - Added agentic retrieval with fallback to standard Retrieve API - search_by_full_text calls Bedrock Retrieve API - Unit tests included - Added README and BEDROCK_MANAGED_KB.md design doc
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for Amazon Bedrock Managed Knowledge Bases and adds agentic retrieval capabilities (query decomposition and managed reranking) with automatic fallback to standard retrieval. The feedback focuses on improving null safety in _get_source_uri and retrieval settings to prevent potential runtime errors, aligning the default knowledge base type between the YAML configuration and Python code, correcting a typo in the required boto3 version, and adding unit tests for the new agentic stream retrieval path.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| def _get_source_uri(result): | ||
| """Extract source URI from a retrieval result, handling all location types.""" | ||
| location = result.get('location', {}) | ||
| loc_type = location.get('type', '') | ||
| if loc_type == 'S3' or 's3Location' in location: | ||
| return location.get('s3Location', {}).get('uri', '') | ||
| elif loc_type == 'WEB' or 'webLocation' in location: | ||
| return location.get('webLocation', {}).get('url', '') | ||
| elif 'confluenceLocation' in location: | ||
| return location.get('confluenceLocation', {}).get('url', '') | ||
| elif 'salesforceLocation' in location: | ||
| return location.get('salesforceLocation', {}).get('url', '') | ||
| elif 'sharePointLocation' in location: | ||
| return location.get('sharePointLocation', {}).get('url', '') | ||
| elif 'customDocumentLocation' in location: | ||
| return location.get('customDocumentLocation', {}).get('id', '') | ||
| # Fallback to metadata._source_uri (for agentic results) | ||
| return result.get('metadata', {}).get('_source_uri', '') |
There was a problem hiding this comment.
In _get_source_uri, if result contains 'location' or 'metadata' with an explicit None value, calling .get() on them will raise an AttributeError. Using result.get('location') or {} and result.get('metadata') or {} ensures robust null safety.
| def _get_source_uri(result): | |
| """Extract source URI from a retrieval result, handling all location types.""" | |
| location = result.get('location', {}) | |
| loc_type = location.get('type', '') | |
| if loc_type == 'S3' or 's3Location' in location: | |
| return location.get('s3Location', {}).get('uri', '') | |
| elif loc_type == 'WEB' or 'webLocation' in location: | |
| return location.get('webLocation', {}).get('url', '') | |
| elif 'confluenceLocation' in location: | |
| return location.get('confluenceLocation', {}).get('url', '') | |
| elif 'salesforceLocation' in location: | |
| return location.get('salesforceLocation', {}).get('url', '') | |
| elif 'sharePointLocation' in location: | |
| return location.get('sharePointLocation', {}).get('url', '') | |
| elif 'customDocumentLocation' in location: | |
| return location.get('customDocumentLocation', {}).get('id', '') | |
| # Fallback to metadata._source_uri (for agentic results) | |
| return result.get('metadata', {}).get('_source_uri', '') | |
| def _get_source_uri(result): | |
| """Extract source URI from a retrieval result, handling all location types.""" | |
| location = result.get('location') or {} | |
| loc_type = location.get('type', '') | |
| if loc_type == 'S3' or 's3Location' in location: | |
| return (location.get('s3Location') or {}).get('uri', '') | |
| elif loc_type == 'WEB' or 'webLocation' in location: | |
| return (location.get('webLocation') or {}).get('url', '') | |
| elif 'confluenceLocation' in location: | |
| return (location.get('confluenceLocation') or {}).get('url', '') | |
| elif 'salesforceLocation' in location: | |
| return (location.get('salesforceLocation') or {}).get('url', '') | |
| elif 'sharePointLocation' in location: | |
| return (location.get('sharePointLocation') or {}).get('url', '') | |
| elif 'customDocumentLocation' in location: | |
| return (location.get('customDocumentLocation') or {}).get('id', '') | |
| # Fallback to metadata._source_uri (for agentic results) | |
| return (result.get('metadata') or {}).get('_source_uri', '') |
| score = retrieval_result.get("score", 1.0) | ||
| if score < retrieval_setting.get("score_threshold", 0.0): | ||
| continue | ||
| result = { | ||
| "metadata": retrieval_result.get("metadata", {}), | ||
| "score": score, | ||
| "title": _get_source_uri(retrieval_result), | ||
| "content": retrieval_result.get("content", {}).get("text", ""), | ||
| } |
There was a problem hiding this comment.
Safeguard against None values for retrieval_setting, score_threshold, metadata, and content to prevent potential TypeError or AttributeError during agentic retrieval.
| score = retrieval_result.get("score", 1.0) | |
| if score < retrieval_setting.get("score_threshold", 0.0): | |
| continue | |
| result = { | |
| "metadata": retrieval_result.get("metadata", {}), | |
| "score": score, | |
| "title": _get_source_uri(retrieval_result), | |
| "content": retrieval_result.get("content", {}).get("text", ""), | |
| } | |
| score = retrieval_result.get("score", 1.0) | |
| threshold = (retrieval_setting or {}).get("score_threshold") or 0.0 | |
| if score < threshold: | |
| continue | |
| result = { | |
| "metadata": retrieval_result.get("metadata") or {}, | |
| "score": score, | |
| "title": _get_source_uri(retrieval_result), | |
| "content": (retrieval_result.get("content") or {}).get("text", ""), | |
| } |
| # Determine retrieval configuration based on knowledge base type | ||
| kb_type = settings.get("knowledge_base_type", "MANAGED") | ||
| if kb_type == "MANAGED": | ||
| retrieval_config = { | ||
| "managedSearchConfiguration": {"numberOfResults": retrieval_setting.get("top_k")} | ||
| } | ||
| else: | ||
| retrieval_config = { | ||
| "vectorSearchConfiguration": {"numberOfResults": retrieval_setting.get("top_k"), | ||
| "overrideSearchType": "HYBRID"} | ||
| }, | ||
| } |
There was a problem hiding this comment.
Safely retrieve top_k with a fallback default value to prevent AttributeError if retrieval_setting is None or missing.
| # Determine retrieval configuration based on knowledge base type | |
| kb_type = settings.get("knowledge_base_type", "MANAGED") | |
| if kb_type == "MANAGED": | |
| retrieval_config = { | |
| "managedSearchConfiguration": {"numberOfResults": retrieval_setting.get("top_k")} | |
| } | |
| else: | |
| retrieval_config = { | |
| "vectorSearchConfiguration": {"numberOfResults": retrieval_setting.get("top_k"), | |
| "overrideSearchType": "HYBRID"} | |
| }, | |
| } | |
| # Determine retrieval configuration based on knowledge base type | |
| kb_type = settings.get("knowledge_base_type", "MANAGED") | |
| top_k = (retrieval_setting or {}).get("top_k", 5) | |
| if kb_type == "MANAGED": | |
| retrieval_config = { | |
| "managedSearchConfiguration": {"numberOfResults": top_k} | |
| } | |
| else: | |
| retrieval_config = { | |
| "vectorSearchConfiguration": {"numberOfResults": top_k, | |
| "overrideSearchType": "HYBRID"} | |
| } |
| - name: knowledge_base_type | ||
| type: select | ||
| required: false | ||
| default: VECTOR |
There was a problem hiding this comment.
The default value in the YAML is set to VECTOR, but the Python code fallback defaults to MANAGED. Since MANAGED is labeled as 'Recommended', consider aligning the YAML default value to MANAGED as well.
default: MANAGED|
|
||
| ### Requirements | ||
|
|
||
| - **boto3 >= 1.43** is required for agentic retrieval support. Earlier versions do not include the necessary API parameters. |
There was a problem hiding this comment.
The boto3 version 1.43 does not exist (boto3 is currently in the 1.35.x/1.36.x range). This is likely a typo for 1.34.143 or 1.35.43.
| - **boto3 >= 1.43** is required for agentic retrieval support. Earlier versions do not include the necessary API parameters. | |
| - **boto3 >= 1.34.143** is required for agentic retrieval support. Earlier versions do not include the necessary API parameters. |
| - Configurable via Dify plugin settings UI | ||
|
|
||
| ## SDK Requirements | ||
| - boto3 >= 1.43 |
There was a problem hiding this comment.
The boto3 version 1.43 does not exist. Please correct this to the intended version (e.g., 1.34.143 or 1.35.43).
| - boto3 >= 1.43 | |
| - boto3 >= 1.34.143 |
| response = endpoint._invoke(request, {}, settings) | ||
| result = json.loads(response.data) | ||
| assert len(result["records"]) == 1 | ||
| assert result["records"][0]["content"] == "High score" |
There was a problem hiding this comment.
Consider adding a unit test to cover the new agentic_retrieve_stream functionality, including its event stream parsing and fallback behavior when an exception occurs.
1 participant
feat: add Managed Knowledge Base support to AWS Bedrock KB plugin
Summary
Adds Managed Knowledge Base type support to the existing AWS Bedrock Knowledge Base extension plugin. Users can now select between MANAGED (recommended) and VECTOR knowledge base types via a dropdown in the provider configuration.
Fixes N/A — new feature.
Change Type
Screenshots / Videos
LLM Plugin Checklist
N/A — this is an extension plugin, not an LLM plugin.
Version
versioninmanifest.yaml(not the one undermeta)dify_plugin>=0.3.0,<0.6.0is declared inpyproject.tomland locked inuv.lock(or kept inrequirements.txtfor legacy plugins withoutuv.lock) — SDK docsTesting
Changes Made
managedSearchConfigurationsearch_by_full_textcalls Bedrock Retrieve API