feat: add Bedrock Knowledge Base capability#379
Conversation
Adds a new BedrockKnowledgeBase capability that connects Pydantic AI agents to Amazon Bedrock Managed Knowledge Bases for RAG (retrieval-augmented generation). Features: - Agentic retrieval (multi-step reasoning via AgenticRetrieveStream) - Standard semantic retrieval (Retrieve API) - Direct document ingestion via IngestKnowledgeBaseDocuments (CUSTOM data source) - Inline text, S3 reference, and binary content modes - S3 upload + StartIngestionJob (S3 data source) Tools exposed: - search_knowledge_base: always available - ingest_document: opt-in via include_ingest_tool=True Requires boto3 >= 1.43.2 (not added to pyproject.toml per contribution policy).
📝 WalkthroughWalkthroughAdds a 🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (3)
pydantic_ai_harness/bedrock_kb/_capability.py (1)
54-56: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winConsider validating
data_source_typelikenumber_of_results.Only
'S3'and'CUSTOM'are meaningful (per the toolset's routing in_toolset.py), but nothing rejects a typo or other value at construction time; it silently falls through to the S3 path in the toolset. Since__post_init__already validatesnumber_of_results, extending it to validatedata_source_typewould catch misconfiguration earlier.🔧 Proposed validation
def __post_init__(self) -> None: if self.number_of_results <= 0: raise ValueError(f'number_of_results must be positive, got {self.number_of_results}') + if self.data_source_type.upper() not in ('S3', 'CUSTOM'): + raise ValueError(f"data_source_type must be 'S3' or 'CUSTOM', got {self.data_source_type!r}")Also applies to: 69-71
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pydantic_ai_harness/bedrock_kb/_capability.py` around lines 54 - 56, Extend the existing __post_init__ validation for number_of_results to validate data_source_type as well, accepting only 'S3' or 'CUSTOM' and rejecting any other value at construction time. Keep the existing routing behavior and validation unchanged for valid configurations.pydantic_ai_harness/bedrock_kb/_toolset.py (2)
101-101: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRepeated inline imports and
get_event_loop()instead ofget_running_loop().
asyncio/uuidare imported inline in three different methods instead of once at module scope, andasyncio.get_event_loop()is used inside coroutines that are already running —asyncio.get_running_loop()is the more explicit, non-deprecated choice in that context.Also applies to: 125-125, 130-130, 150-150, 181-182, 192-192
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pydantic_ai_harness/bedrock_kb/_toolset.py` at line 101, Move the repeated asyncio and uuid imports from the affected methods into module scope, then remove the inline imports while preserving their existing usage. In each coroutine currently calling asyncio.get_event_loop(), replace it with asyncio.get_running_loop(), including the flows around the relevant tool methods.
45-46: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
Anyand un-parameterizeddictfor boto3 clients/kwargs.
_runtime_client/_agent_clientare typedAnyandkwargs: dict(line 133) has no type parameters. As per coding guidelines,**/*.pyshould "provide full type annotations, avoidAny, and do not usecast(); use type narrowing instead." Ifboto3-stubs/mypy-boto3-bedrock-agent-runtimearen't already a dependency, at least parameterizedict[str, Any]forkwargs.Also applies to: 52-63, 65-76, 133-133
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pydantic_ai_harness/bedrock_kb/_toolset.py` around lines 45 - 46, Replace the untyped boto3 client fields in the toolset initializer and related methods with precise Bedrock Agent Runtime and Agent client types, using the project’s available stubs or interfaces without introducing cast(). Parameterize the kwargs annotation in the affected method as dict[str, Any] at minimum, and apply the same typing consistently across the referenced client methods.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@pydantic_ai_harness/bedrock_kb/_toolset.py`:
- Around line 128-151: Update _standard_retrieve to honor
self._number_of_results by limiting the retrieval results before formatting or
truncating the formatted results before joining them; preserve the existing
no-results response and output formatting.
- Around line 253-254: Update the missing-bucket guard in the relevant tool
method to raise ModelRetry with the existing “No S3 bucket configured for
ingestion” message instead of returning a plain string. Match the established
ModelRetry handling used by the knowledge_base_id and data_source_id
configuration guards.
- Around line 99-126: The _agentic_retrieve method currently calls
retrieve_and_generate, which is incompatible with managed knowledge bases.
Replace that runtime call with agentic_retrieve_stream, preserving the existing
asynchronous executor flow and adapting response parsing to the stream’s output
and citations while retaining the configured result limit.
- Around line 201-209: The s3_uri branch in ingest_document currently allows
unrestricted S3 ingestion; restrict caller-supplied URIs to the configured
allowlisted bucket and key prefix, and include bucketOwnerAccountId in the
CUSTOM S3 location when required. Reject or avoid forwarding URIs outside that
allowlist, preserving the existing non-S3 ingestion path.
- Around line 212-225: Update ingest_document’s binary-content handling to
base64-decode content into bytes before uploading or embedding it. Apply the
decode once in the BYTE inlineContent branch shown by the mime_type path and
before the S3 put_object call, while leaving non-binary content handling
unchanged.
In `@pydantic_ai_harness/bedrock_kb/README.md`:
- Around line 7-10: Replace the em-dashes separating the bold feature labels
from their descriptions in the README feature list with plain ASCII punctuation,
preserving the existing wording and formatting.
In `@tests/bedrock_kb/test_bedrock_kb.py`:
- Around line 13-27: Update the tests in the Bedrock knowledge-base test module
to use pytest-anyio async test methods instead of synchronous tests that drive
coroutines with asyncio.run(). Replace imports from private _toolset and
_capability modules with imports of BedrockKBToolset and BedrockKnowledgeBase
from the public pydantic_ai_harness.bedrock_kb package, including the
_make_toolset helper and all affected tests through the referenced range.
- Around line 10-131: Extend TestBedrockKBToolset to cover the
use_agentic_retrieval=True path through search_knowledge_base, plus
ingest_document validation for missing knowledge_base_id and missing S3 bucket,
and exception-to-ModelRetry handling for retrieval and ingestion. Mock the
relevant clients and assert the expected API calls, returned results, and
ModelRetry messages so all branches in the BedrockKBToolset flows are exercised.
---
Nitpick comments:
In `@pydantic_ai_harness/bedrock_kb/_capability.py`:
- Around line 54-56: Extend the existing __post_init__ validation for
number_of_results to validate data_source_type as well, accepting only 'S3' or
'CUSTOM' and rejecting any other value at construction time. Keep the existing
routing behavior and validation unchanged for valid configurations.
In `@pydantic_ai_harness/bedrock_kb/_toolset.py`:
- Line 101: Move the repeated asyncio and uuid imports from the affected methods
into module scope, then remove the inline imports while preserving their
existing usage. In each coroutine currently calling asyncio.get_event_loop(),
replace it with asyncio.get_running_loop(), including the flows around the
relevant tool methods.
- Around line 45-46: Replace the untyped boto3 client fields in the toolset
initializer and related methods with precise Bedrock Agent Runtime and Agent
client types, using the project’s available stubs or interfaces without
introducing cast(). Parameterize the kwargs annotation in the affected method as
dict[str, Any] at minimum, and apply the same typing consistently across the
referenced client methods.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 78a550f6-40d3-40df-ba7f-a9361849eb30
📒 Files selected for processing (6)
pydantic_ai_harness/bedrock_kb/README.mdpydantic_ai_harness/bedrock_kb/__init__.pypydantic_ai_harness/bedrock_kb/_capability.pypydantic_ai_harness/bedrock_kb/_toolset.pytests/bedrock_kb/__init__.pytests/bedrock_kb/test_bedrock_kb.py
🔗 Linked repositories identified
CodeRabbit considers these linked repositories for cross-repo context during reviews:
pydantic/logfire(manual)pydantic/pydantic-ai(manual)pydantic/monty(auto-detected)pydantic/pydantic(auto-detected)
| async def _agentic_retrieve(self, query: str) -> str: | ||
| """Use AgenticRetrieveStream for multi-step reasoning retrieval.""" | ||
| import asyncio | ||
|
|
||
| def _call() -> str: | ||
| response = self._get_runtime_client.retrieve_and_generate( | ||
| input={'text': query}, | ||
| retrieveAndGenerateConfiguration={ | ||
| 'type': 'KNOWLEDGE_BASE', | ||
| 'knowledgeBaseConfiguration': { | ||
| 'knowledgeBaseId': self._knowledge_base_id, | ||
| 'modelArn': 'arn:aws:bedrock:us-west-2::foundation-model/anthropic.claude-3-haiku-20240307-v1:0', | ||
| }, | ||
| }, | ||
| ) | ||
| # RetrieveAndGenerate returns generated output + citations | ||
| output = response.get('output', {}).get('text', '') | ||
| citations = response.get('citations', []) | ||
| results = [f'Answer: {output}'] | ||
| for citation in citations[:self._number_of_results]: | ||
| for ref in citation.get('retrievedReferences', []): | ||
| content = ref.get('content', {}).get('text', '') | ||
| source = ref.get('location', {}).get('s3Location', {}).get('uri', 'unknown') | ||
| results.append(f'(Source: {source})\n{content[:200]}...') | ||
| return '\n\n---\n\n'.join(results) if results else 'No results found.' | ||
|
|
||
| loop = asyncio.get_event_loop() | ||
| return await loop.run_in_executor(None, _call) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | 🏗️ Heavy lift
🧩 Analysis chain
🌐 Web query:
Does boto3 bedrock-agent-runtime agentic_retrieve_stream require different request/response shape than retrieve_and_generate?
💡 Result:
Yes, agentic_retrieve_stream and retrieve_and_generate require different request and response structures because they serve different architectural purposes and operate on different types of knowledge base configurations [1][2]. Request Shape The retrieve_and_generate operation is designed for standard knowledge bases and requires an input structure that specifies the knowledge base ID and configuration (such as the model ARN and retrieval settings) [3][4]. In contrast, agentic_retrieve_stream is specifically used for agentic retrieval, which utilizes a foundation model to intelligently decompose complex queries [5]. Because this process is agentic, the request parameters are focused on the agentic context rather than the rigid, static configuration required by the standard retrieve_and_generate API [2]. Response Shape The response structures also differ significantly: 1. retrieve_and_generate: This is a standard request-response operation that returns a complete, non-streamed object containing the retrieved results and the generated response in a single payload [6][4]. 2. agentic_retrieve_stream: As the name implies, this operation returns a stream [7][2]. The response consists of a series of events, including: - Result events (AgenticRetrieveResultEvent): These contain the final retrieval results and, if response generation is enabled, the synthesized answer and citations [2]. - Response events (AgenticRetrieveResponseEvent): These are streamed incrementally during the generation process, providing portions of the natural-language answer as they are produced [2]. - Trace events: These are provided throughout the stream to offer visibility into the agentic decomposition and retrieval process [2]. Due to these differences, you cannot use retrieve_and_generate for managed knowledge bases that require agentic retrieval, nor can you swap the invocation methods interchangeably [1][4].
Citations:
- 1: https://docs.aws.amazon.com/boto3/latest/reference/services/bedrock-agent-runtime/client/retrieve_and_generate.html
- 2: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-agentic-retrieve.html
- 3: https://boto3.amazonaws.com/v1/documentation/api/1.34.67/reference/services/bedrock-agent-runtime/client/retrieve_and_generate.html
- 4: https://docs.aws.amazon.com/cli/latest/reference/bedrock-agent-runtime/retrieve-and-generate.html
- 5: https://docs.aws.amazon.com/boto3/latest/reference/services/bedrock-agent-runtime/client/agentic_retrieve_stream.html
- 6: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_RetrieveAndGenerate.html
- 7: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_AgenticRetrieveStream.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## file outline\n'
ast-grep outline pydantic_ai_harness/bedrock_kb/_toolset.py --view expanded || true
printf '\n## relevant file slice\n'
sed -n '1,180p' pydantic_ai_harness/bedrock_kb/_toolset.py
printf '\n## related references\n'
rg -n "AgenticRetrieveStream|retrieve_and_generate|managed knowledge|use_agentic_retrieval|number_of_results" pydantic_ai_harness/bedrock_kb -SRepository: pydantic/pydantic-ai-harness
Length of output: 10268
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## README and capability docs mentions\n'
rg -n "AgenticRetrieveStream|retrieve_and_generate|managed knowledge|use_agentic_retrieval" README.md pydantic_ai_harness -S || trueRepository: pydantic/pydantic-ai-harness
Length of output: 1596
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## file outline\n'
ast-grep outline pydantic_ai_harness/bedrock_kb/_toolset.py --view expanded || true
printf '\n## relevant file slice\n'
sed -n '1,180p' pydantic_ai_harness/bedrock_kb/_toolset.py
printf '\n## related references\n'
rg -n "AgenticRetrieveStream|retrieve_and_generate|managed knowledge|use_agentic_retrieval|number_of_results" pydantic_ai_harness/bedrock_kb -S || trueRepository: pydantic/pydantic-ai-harness
Length of output: 10268
Use agentic_retrieve_stream here. retrieve_and_generate does not support managed knowledge bases, so the default agentic search path will fail. The docstring and README already point to AgenticRetrieveStream.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pydantic_ai_harness/bedrock_kb/_toolset.py` around lines 99 - 126, The
_agentic_retrieve method currently calls retrieve_and_generate, which is
incompatible with managed knowledge bases. Replace that runtime call with
agentic_retrieve_stream, preserving the existing asynchronous executor flow and
adapting response parsing to the stream’s output and citations while retaining
the configured result limit.
| async def _standard_retrieve(self, query: str) -> str: | ||
| """Use standard Retrieve API for semantic search.""" | ||
| import asyncio | ||
|
|
||
| def _call() -> str: | ||
| kwargs: dict = { | ||
| 'knowledgeBaseId': self._knowledge_base_id, | ||
| 'retrievalQuery': {'text': query}, | ||
| } | ||
| # managedSearchConfiguration for managed KBs, vectorSearchConfiguration for vector KBs | ||
| # Let the API auto-detect by not specifying configuration (works for both) | ||
| response = self._get_runtime_client.retrieve(**kwargs) | ||
| results = [] | ||
| for r in response.get('retrievalResults', []): | ||
| content = r.get('content', {}).get('text', '') | ||
| score = r.get('score', 0) | ||
| source = r.get('location', {}).get('s3Location', {}).get('uri', 'unknown') | ||
| results.append(f'[Score: {score:.2f}] (Source: {source})\n{content}') | ||
| if not results: | ||
| return 'No results found.' | ||
| return '\n\n---\n\n'.join(results) | ||
|
|
||
| loop = asyncio.get_event_loop() | ||
| return await loop.run_in_executor(None, _call) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
number_of_results is ignored by standard retrieval.
_standard_retrieve never sets a result-limiting configuration on the retrieve() call and never truncates response['retrievalResults'], so self._number_of_results (documented as "Maximum number of results to return from retrieval") has no effect when use_agentic_retrieval=False. At minimum, slice the formatted results to self._number_of_results.
🔧 Proposed fix
response = self._get_runtime_client.retrieve(**kwargs)
results = []
- for r in response.get('retrievalResults', []):
+ for r in response.get('retrievalResults', [])[: self._number_of_results]:
content = r.get('content', {}).get('text', '')📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| async def _standard_retrieve(self, query: str) -> str: | |
| """Use standard Retrieve API for semantic search.""" | |
| import asyncio | |
| def _call() -> str: | |
| kwargs: dict = { | |
| 'knowledgeBaseId': self._knowledge_base_id, | |
| 'retrievalQuery': {'text': query}, | |
| } | |
| # managedSearchConfiguration for managed KBs, vectorSearchConfiguration for vector KBs | |
| # Let the API auto-detect by not specifying configuration (works for both) | |
| response = self._get_runtime_client.retrieve(**kwargs) | |
| results = [] | |
| for r in response.get('retrievalResults', []): | |
| content = r.get('content', {}).get('text', '') | |
| score = r.get('score', 0) | |
| source = r.get('location', {}).get('s3Location', {}).get('uri', 'unknown') | |
| results.append(f'[Score: {score:.2f}] (Source: {source})\n{content}') | |
| if not results: | |
| return 'No results found.' | |
| return '\n\n---\n\n'.join(results) | |
| loop = asyncio.get_event_loop() | |
| return await loop.run_in_executor(None, _call) | |
| async def _standard_retrieve(self, query: str) -> str: | |
| """Use standard Retrieve API for semantic search.""" | |
| import asyncio | |
| def _call() -> str: | |
| kwargs: dict = { | |
| 'knowledgeBaseId': self._knowledge_base_id, | |
| 'retrievalQuery': {'text': query}, | |
| } | |
| # managedSearchConfiguration for managed KBs, vectorSearchConfiguration for vector KBs | |
| # Let the API auto-detect by not specifying configuration (works for both) | |
| response = self._get_runtime_client.retrieve(**kwargs) | |
| results = [] | |
| for r in response.get('retrievalResults', [])[: self._number_of_results]: | |
| content = r.get('content', {}).get('text', '') | |
| score = r.get('score', 0) | |
| source = r.get('location', {}).get('s3Location', {}).get('uri', 'unknown') | |
| results.append(f'[Score: {score:.2f}] (Source: {source})\n{content}') | |
| if not results: | |
| return 'No results found.' | |
| return '\n\n---\n\n'.join(results) | |
| loop = asyncio.get_event_loop() | |
| return await loop.run_in_executor(None, _call) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pydantic_ai_harness/bedrock_kb/_toolset.py` around lines 128 - 151, Update
_standard_retrieve to honor self._number_of_results by limiting the retrieval
results before formatting or truncating the formatted results before joining
them; preserve the existing no-results response and output formatting.
| if s3_uri: | ||
| doc = { | ||
| 'content': { | ||
| 'dataSourceType': 'CUSTOM', | ||
| 'custom': { | ||
| 'customDocumentIdentifier': {'id': doc_id}, | ||
| 'sourceType': 'S3_LOCATION', | ||
| 's3Location': {'uri': s3_uri}, | ||
| }, |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## file outline\n'
ast-grep outline pydantic_ai_harness/bedrock_kb/_toolset.py --view expanded || true
printf '\n## relevant lines\n'
sed -n '1,280p' pydantic_ai_harness/bedrock_kb/_toolset.pyRepository: pydantic/pydantic-ai-harness
Length of output: 11505
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## search for s3_uri usage\n'
rg -n "s3_uri|bucketOwnerAccountId|search_knowledge_base|ingestion" pydantic_ai_harness -SRepository: pydantic/pydantic-ai-harness
Length of output: 3592
🌐 Web query:
Amazon Bedrock knowledge base ingest_knowledge_base_documents s3Location service role bucketOwnerAccountId
💡 Result:
The parameter bucketOwnerAccountId in the Amazon Bedrock ingest_knowledge_base_documents API operation is a string that specifies the AWS account ID of the owner of the S3 bucket containing the documents to be ingested [1][2]. Key details regarding its use include: 1. Function: It is primarily used to identify the owner of the S3 bucket when the bucket resides in a different AWS account than the one managing the Bedrock knowledge base (cross-account access) [3][4]. 2. Requirements for Cross-Account Access: If the S3 bucket belongs to a different AWS account, you must ensure the Bedrock knowledge base service role has the appropriate permissions to access that bucket [3][4]. This typically involves configuring the S3 bucket policy in the owning account to grant access to the service role ARN from your knowledge base account [3][4]. 3. Service Role Configuration: Your Amazon Bedrock Knowledge Base service role must have permissions to access the specified S3 bucket [3]. If the bucket is cross-account or encrypted with a customer-managed KMS key, the service role policy (and the bucket/KMS policies) must be explicitly updated to include the bucketOwnerAccountId and the necessary permissions for the target resources [3]. When using the ingest_knowledge_base_documents operation, s3Location is a required dictionary parameter that defines where the documents are located, and within that structure, bucketOwnerAccountId helps Amazon Bedrock correctly resolve the S3 object path [1][5].
Citations:
- 1: https://docs.aws.amazon.com/boto3/latest/reference/services/bedrock-agent/client/ingest_knowledge_base_documents.html
- 2: https://docs.aws.amazon.com/cli/latest/reference/bedrock-agent/ingest-knowledge-base-documents.html
- 3: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-permissions.html
- 4: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-managed-ds-s3.html
- 5: https://boto3.amazonaws.com/v1/documentation/api/1.40.0/reference/services/bedrock-agent/client/ingest_knowledge_base_documents.html
SSRF (CWE-918): Server-Side Request Forgery (SSRF)
Reachability: External
● Entry
pydantic_ai_harness/bedrock_kb/__init__.py:3
BedrockKnowledgeBase
│
▼
● Hop
pydantic_ai_harness/bedrock_kb/_capability.py:73
get_toolset: Build and return the Bedrock KB toolset.
│
▼
● Sink
pydantic_ai_harness/bedrock_kb/_toolset.py
Restrict caller-supplied s3_uri in the CUSTOM ingestion path
ingest_document(..., s3_uri=...) forwards the URI straight to Bedrock, so a tool caller can make the KB service role ingest any S3 object it already can read and surface it through search. Constrain this to an allowlisted bucket and key prefix, set bucketOwnerAccountId, or remove the S3-reference branch for model calls.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pydantic_ai_harness/bedrock_kb/_toolset.py` around lines 201 - 209, The
s3_uri branch in ingest_document currently allows unrestricted S3 ingestion;
restrict caller-supplied URIs to the configured allowlisted bucket and key
prefix, and include bucketOwnerAccountId in the CUSTOM S3 location when
required. Reject or avoid forwarding URIs outside that allowlist, preserving the
existing non-S3 ingestion path.
| elif mime_type: | ||
| doc = { | ||
| 'content': { | ||
| 'dataSourceType': 'CUSTOM', | ||
| 'custom': { | ||
| 'customDocumentIdentifier': {'id': doc_id}, | ||
| 'sourceType': 'IN_LINE', | ||
| 'inlineContent': { | ||
| 'type': 'BYTE', | ||
| 'byteContent': {'data': content, 'mimeType': mime_type}, | ||
| }, | ||
| }, | ||
| }, | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- file outline ---'
python3 - <<'PY'
from pathlib import Path
p = Path('pydantic_ai_harness/bedrock_kb/_toolset.py')
for i, line in enumerate(p.read_text().splitlines(), 1):
if 180 <= i <= 290:
print(f'{i:4}: {line}')
PY
echo '--- search base64/content usage ---'
rg -n "base64|content" pydantic_ai_harness/bedrock_kb/_toolset.py pydantic_ai_harness/bedrock_kb -g '!**/__pycache__/**'Repository: pydantic/pydantic-ai-harness
Length of output: 9331
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path('pydantic_ai_harness/bedrock_kb/_toolset.py')
for i, line in enumerate(p.read_text().splitlines(), 1):
if 180 <= i <= 290:
print(f'{i:4}: {line}')
PY
echo '---'
rg -n "base64|content" pydantic_ai_harness/bedrock_kb/_toolset.py pydantic_ai_harness/bedrock_kb -g '!**/__pycache__/**'Repository: pydantic/pydantic-ai-harness
Length of output: 9278
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- relevant README section ---'
python3 - <<'PY'
from pathlib import Path
p = Path('pydantic_ai_harness/bedrock_kb/README.md')
for i, line in enumerate(p.read_text().splitlines(), 1):
if 45 <= i <= 75:
print(f'{i:4}: {line}')
PY
echo '--- tests/search ---'
rg -n "ingest_document|mime_type|base64|BYTE|put_object|start_ingestion_job" . -g '!**/__pycache__/**' -g '!**/.git/**'Repository: pydantic/pydantic-ai-harness
Length of output: 11288
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path('tests/bedrock_kb/test_bedrock_kb.py')
for i, line in enumerate(p.read_text().splitlines(), 1):
if 1 <= i <= 220:
print(f'{i:4}: {line}')
PYRepository: pydantic/pydantic-ai-harness
Length of output: 7114
🌐 Web query:
botocore blob field string bytes base64 serialization boto3 blob parameter str encoded
💡 Result:
In Boto3 and Botocore, parameters modeled as blob types are designed to handle binary data. While official documentation typically specifies bytes or file-like objects as the required input, the underlying serialization behavior involves specific processing [1][2][3]. Key aspects of blob serialization: 1. Handling of Data Types: Botocore validates blob parameters to accept bytes, bytearray, or file-like objects (which must implement a read method and return bytes) [1][4]. While passing a str (string) sometimes works due to legacy Python 2.7 compatibility quirks, it is generally discouraged, not officially supported, and can lead to runtime errors or TypeErrors in strict environments (such as those using type checkers like mypy) [3][5][6][7]. 2. Serialization Process: When you provide bytes to a blob parameter, Botocore does not perform manual base64 encoding before transit; rather, the serialization layer processes the input, and the underlying request handling ensures the data is transmitted correctly to the AWS service as binary data [8][9][3]. In many AWS service API models, blob fields are expected to be raw bytes, and the SDK manages the transition for the wire protocol [3]. 3. Best Practices: To ensure compatibility and avoid errors: Always encode string data (e.g., JSON payloads) into bytes before passing it to a blob parameter [6]. Example: data_bytes = my_string.encode("utf-8") [6]. For file uploads, always open files in binary mode (mode="rb") and pass the file-like object directly [10][7][4]. Do not rely on the SDK to automatically handle string-to-byte encoding, as this behavior is considered a legacy implementation detail rather than a robust, intended feature [3][5]. For specific services like DynamoDB, use the specialized boto3.dynamodb.types.Binary wrapper when explicit binary representation is required in item definitions [11]. For general client operations, sticking to bytes or binary file-like objects is the standard [3][6].
Citations:
- 1: https://github.com/boto/botocore/blob/a748cc02/botocore/validate.py
- 2: Docs: S3 put_object Body parameter can be a text object boto/boto3#2407
- 3: Docs for iot-data publish() method are incorrect boto/botocore#2619
- 4: https://docs.aws.amazon.com/boto3/latest/reference/services/s3/bucket/upload_fileobj.html
- 5: Docs: lambda.invoke Payload arg also takes type str boto/boto3#1543
- 6: https://brandonwie.dev/posts/boto3-s3-encoding
- 7: Type error when submitting to S3 bucket boto/boto3#1111
- 8: https://github.com/boto/botocore/blob/develop/botocore/serialize.py
- 9: https://github.com/boto/botocore/blob/a748cc02/tests/unit/test_serialize.py
- 10: https://docs.aws.amazon.com/boto3/latest/guide/s3-uploading-files.html
- 11: https://docs.aws.amazon.com/boto3/latest/reference/customizations/dynamodb.html
🌐 Web query:
AWS Bedrock ingest_knowledge_base_documents byteContent data blob base64 boto3
💡 Result:
To ingest documents directly into an Amazon Bedrock knowledge base using the boto3 ingest_knowledge_base_documents method, you provide document content via the byteContent field when the type is set to BYTE [1][2]. Although the official documentation often describes the data field within byteContent as requiring a "base64-encoded string," the Boto3 library accepts this field as bytes (e.g., b'your-content-here') [1][2]. When you pass the data as bytes to the Boto3 client, the underlying botocore library handles the necessary Base64 encoding automatically before sending the request to the Amazon Bedrock API [1]. Key Implementation Details: 1. Data Structure: The byteContent dictionary requires a mimeType (string) and the data field (bytes) [1][2]. 2. Configuration: Ensure your custom content configuration specifies sourceType: 'IN_LINE' and type: 'BYTE' [1][2]. 3. Example snippet: python import boto3 client = boto3.client('bedrock-agent') response = client.ingest_knowledge_base_documents( knowledgeBaseId='your-kb-id', dataSourceId='your-data-source-id', documents=[ { 'content': { 'dataSourceType': 'CUSTOM', 'custom': { 'customDocumentIdentifier': {'id': 'doc-1'}, 'sourceType': 'IN_LINE', 'inlineContent': { 'type': 'BYTE', 'byteContent': { 'mimeType': 'application/pdf', # or appropriate mimeType 'data': b'your-binary-data-here' # Pass as bytes } } } } } ] ) The documentation for the AWS CLI or lower-level API references often explicitly mention "base64-encoded string" because those interfaces operate at a level where the data must be serialized as a string [3][4]. However, when using Boto3, passing a bytes object is the standard and expected way to satisfy this requirement [1][2].
Citations:
- 1: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent/client/ingest_knowledge_base_documents.html
- 2: https://boto3.amazonaws.com/v1/documentation/api/1.40.0/reference/services/bedrock-agent/client/ingest_knowledge_base_documents.html
- 3: https://docs.aws.amazon.com/boto3/latest/reference/services/bedrock-agent/client/ingest_knowledge_base_documents.html
- 4: https://docs.aws.amazon.com/cli/latest/reference/bedrock-agent/ingest-knowledge-base-documents.html
Decode binary content before upload
ingest_document(..., mime_type=...) documents content as base64 text, but both binary paths forward that encoded string unchanged. That stores/sends base64 text instead of the original bytes. Decode the payload once in the BYTE branch and before the S3 put_object call.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pydantic_ai_harness/bedrock_kb/_toolset.py` around lines 212 - 225, Update
ingest_document’s binary-content handling to base64-decode content into bytes
before uploading or embedding it. Apply the decode once in the BYTE
inlineContent branch shown by the mime_type path and before the S3 put_object
call, while leaving non-binary content handling unchanged.
| if not self._data_source_bucket: | ||
| return 'Error: No S3 bucket configured for ingestion.' |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Missing S3 bucket returns a plain string instead of ModelRetry.
Every other configuration guard in this file (knowledge_base_id, data_source_id) raises ModelRetry so the agent knows to reconfigure/retry. This one silently returns 'Error: ...' as a normal tool result, so the model won't get the same retry signal and may treat it as a successful response.
🔧 Proposed fix
if not self._data_source_bucket:
- return 'Error: No S3 bucket configured for ingestion.'
+ raise ModelRetry('No S3 bucket configured for ingestion. Set BEDROCK_DATA_SOURCE_BUCKET.')📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if not self._data_source_bucket: | |
| return 'Error: No S3 bucket configured for ingestion.' | |
| if not self._data_source_bucket: | |
| raise ModelRetry('No S3 bucket configured for ingestion. Set BEDROCK_DATA_SOURCE_BUCKET.') |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pydantic_ai_harness/bedrock_kb/_toolset.py` around lines 253 - 254, Update
the missing-bucket guard in the relevant tool method to raise ModelRetry with
the existing “No S3 bucket configured for ingestion” message instead of
returning a plain string. Match the established ModelRetry handling used by the
knowledge_base_id and data_source_id configuration guards.
| - **Agentic retrieval** — Multi-step, reasoning-enhanced search via `AgenticRetrieveStream` | ||
| - **Standard retrieval** — Semantic vector search via `Retrieve` | ||
| - **Direct ingestion** — Add documents without S3 using `IngestKnowledgeBaseDocuments` (CUSTOM data source) | ||
| - **S3 ingestion** — Upload to S3 + trigger sync (S3 data source) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Replace em-dashes with ASCII punctuation.
Lines 7-10 use "—" between the bold lead-in and description. As per coding guidelines, "Do not use em-dashes; use -- or separate sentences, and prefer plain ASCII punctuation in prose and comments."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pydantic_ai_harness/bedrock_kb/README.md` around lines 7 - 10, Replace the
em-dashes separating the bold feature labels from their descriptions in the
README feature list with plain ASCII punctuation, preserving the existing
wording and formatting.
Source: Coding guidelines
| class TestBedrockKBToolset: | ||
| """Tests for BedrockKBToolset.""" | ||
|
|
||
| def _make_toolset(self, **kwargs): | ||
| from pydantic_ai_harness.bedrock_kb._toolset import BedrockKBToolset | ||
|
|
||
| defaults = { | ||
| 'knowledge_base_id': 'TEST_KB', | ||
| 'region_name': 'us-west-2', | ||
| 'data_source_id': 'TEST_DS', | ||
| 'data_source_type': 'CUSTOM', | ||
| 'data_source_bucket': '', | ||
| 'use_agentic_retrieval': False, | ||
| 'number_of_results': 5, | ||
| 'include_ingest_tool': True, | ||
| } | ||
| defaults.update(kwargs) | ||
| return BedrockKBToolset(**defaults) | ||
|
|
||
| def test_standard_retrieve(self): | ||
| ts = self._make_toolset() | ||
| mock_client = MagicMock() | ||
| mock_client.retrieve.return_value = { | ||
| 'retrievalResults': [ | ||
| {'content': {'text': 'Result 1'}, 'score': 0.95, 'location': {'s3Location': {'uri': 's3://b/k'}}}, | ||
| {'content': {'text': 'Result 2'}, 'score': 0.80, 'location': {'s3Location': {'uri': 's3://b/k2'}}}, | ||
| ] | ||
| } | ||
| ts._runtime_client = mock_client | ||
|
|
||
| import asyncio | ||
| result = asyncio.run(ts.search_knowledge_base('test query')) | ||
|
|
||
| assert 'Result 1' in result | ||
| assert 'Result 2' in result | ||
| assert '0.95' in result | ||
| mock_client.retrieve.assert_called_once() | ||
|
|
||
| def test_ingest_direct_inline_text(self): | ||
| ts = self._make_toolset() | ||
| mock_client = MagicMock() | ||
| mock_client.ingest_knowledge_base_documents.return_value = { | ||
| 'documentDetails': [{'status': 'STARTING'}] | ||
| } | ||
| ts._agent_client = mock_client | ||
|
|
||
| import asyncio | ||
| result = asyncio.run(ts.ingest_document(content='Hello world', document_id='doc-001')) | ||
|
|
||
| assert 'doc-001' in result | ||
| assert 'STARTING' in result | ||
| call_kwargs = mock_client.ingest_knowledge_base_documents.call_args.kwargs | ||
| doc = call_kwargs['documents'][0] | ||
| assert doc['content']['custom']['inlineContent']['type'] == 'TEXT' | ||
| assert doc['content']['custom']['inlineContent']['textContent']['data'] == 'Hello world' | ||
|
|
||
| def test_ingest_direct_s3_reference(self): | ||
| ts = self._make_toolset() | ||
| mock_client = MagicMock() | ||
| mock_client.ingest_knowledge_base_documents.return_value = { | ||
| 'documentDetails': [{'status': 'STARTING'}] | ||
| } | ||
| ts._agent_client = mock_client | ||
|
|
||
| import asyncio | ||
| result = asyncio.run(ts.ingest_document(content='', s3_uri='s3://bucket/file.pdf', document_id='s3-001')) | ||
|
|
||
| assert 's3-001' in result | ||
| doc = mock_client.ingest_knowledge_base_documents.call_args.kwargs['documents'][0] | ||
| assert doc['content']['custom']['sourceType'] == 'S3_LOCATION' | ||
| assert doc['content']['custom']['s3Location']['uri'] == 's3://bucket/file.pdf' | ||
|
|
||
| def test_ingest_direct_binary(self): | ||
| ts = self._make_toolset() | ||
| mock_client = MagicMock() | ||
| mock_client.ingest_knowledge_base_documents.return_value = { | ||
| 'documentDetails': [{'status': 'STARTING'}] | ||
| } | ||
| ts._agent_client = mock_client | ||
|
|
||
| import asyncio | ||
| result = asyncio.run(ts.ingest_document(content='base64data', mime_type='application/pdf', document_id='bin-001')) | ||
|
|
||
| assert 'bin-001' in result | ||
| doc = mock_client.ingest_knowledge_base_documents.call_args.kwargs['documents'][0] | ||
| assert doc['content']['custom']['inlineContent']['type'] == 'BYTE' | ||
| assert doc['content']['custom']['inlineContent']['byteContent']['mimeType'] == 'application/pdf' | ||
|
|
||
| def test_ingest_s3_mode(self): | ||
| ts = self._make_toolset(data_source_type='S3', data_source_bucket='test-bucket') | ||
| mock_agent = MagicMock() | ||
| ts._agent_client = mock_agent | ||
|
|
||
| with patch('boto3.client') as mock_boto: | ||
| mock_s3 = MagicMock() | ||
| mock_boto.return_value = mock_s3 | ||
|
|
||
| import asyncio | ||
| result = asyncio.run(ts.ingest_document(content='Doc content', document_id='s3-doc')) | ||
|
|
||
| assert 's3-doc' in result | ||
| assert 'uploaded' in result | ||
| mock_agent.start_ingestion_job.assert_called_once() | ||
|
|
||
| def test_no_kb_id_raises_model_retry(self): | ||
| from pydantic_ai.exceptions import ModelRetry | ||
|
|
||
| ts = self._make_toolset(knowledge_base_id='') | ||
|
|
||
| import asyncio | ||
| with pytest.raises(ModelRetry, match='Knowledge Base ID'): | ||
| asyncio.run(ts.search_knowledge_base('test')) | ||
|
|
||
| def test_no_ds_id_raises_model_retry_on_ingest(self): | ||
| from pydantic_ai.exceptions import ModelRetry | ||
|
|
||
| ts = self._make_toolset(data_source_id='') | ||
|
|
||
| import asyncio | ||
| with pytest.raises(ModelRetry, match='Data source ID'): | ||
| asyncio.run(ts.ingest_document(content='test')) | ||
|
|
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Branch coverage gaps, including the default agentic-retrieval path.
No test drives use_agentic_retrieval=True, so _agentic_retrieve (where the critical API-mismatch bug lives) is entirely untested; the missing-KB-id branch of ingest_document, the missing-S3-bucket branch, and the exception→ModelRetry branches are also uncovered. As per coding guidelines, this repo must "Maintain 100% branch coverage, enforced by make testcov."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/bedrock_kb/test_bedrock_kb.py` around lines 10 - 131, Extend
TestBedrockKBToolset to cover the use_agentic_retrieval=True path through
search_knowledge_base, plus ingest_document validation for missing
knowledge_base_id and missing S3 bucket, and exception-to-ModelRetry handling
for retrieval and ingestion. Mock the relevant clients and assert the expected
API calls, returned results, and ModelRetry messages so all branches in the
BedrockKBToolset flows are exercised.
Source: Coding guidelines
| def _make_toolset(self, **kwargs): | ||
| from pydantic_ai_harness.bedrock_kb._toolset import BedrockKBToolset | ||
|
|
||
| defaults = { | ||
| 'knowledge_base_id': 'TEST_KB', | ||
| 'region_name': 'us-west-2', | ||
| 'data_source_id': 'TEST_DS', | ||
| 'data_source_type': 'CUSTOM', | ||
| 'data_source_bucket': '', | ||
| 'use_agentic_retrieval': False, | ||
| 'number_of_results': 5, | ||
| 'include_ingest_tool': True, | ||
| } | ||
| defaults.update(kwargs) | ||
| return BedrockKBToolset(**defaults) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Tests bypass pytest-anyio and import from private modules.
Two guideline gaps here: async behavior is driven with raw asyncio.run() inside sync def test_... methods rather than pytest-anyio, and BedrockKBToolset/BedrockKnowledgeBase are imported from the private _toolset/_capability modules instead of the public pydantic_ai_harness.bedrock_kb package that __init__.py already re-exports them from. As per coding guidelines, "Use pytest-anyio for asynchronous tests" and "Exercise behavior through public capability APIs or public classes re-exported by the capability package."
Also applies to: 29-131
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/bedrock_kb/test_bedrock_kb.py` around lines 13 - 27, Update the tests
in the Bedrock knowledge-base test module to use pytest-anyio async test methods
instead of synchronous tests that drive coroutines with asyncio.run(). Replace
imports from private _toolset and _capability modules with imports of
BedrockKBToolset and BedrockKnowledgeBase from the public
pydantic_ai_harness.bedrock_kb package, including the _make_toolset helper and
all affected tests through the referenced range.
Source: Coding guidelines
1 participant
Summary
Adds a new BedrockKnowledgeBase capability that connects Pydantic AI agents to Amazon Bedrock Managed Knowledge Bases for RAG.
Features:
Tools exposed:
Requires boto3 >= 1.43.2 (not added to pyproject.toml per contribution policy).
Linked Issue
Fixes #378
Checklist