From 77b4c1832405c6f784ce67233ad919fddaa6e91a Mon Sep 17 00:00:00 2001 From: PVidyadhar Date: Thu, 16 Jul 2026 09:35:49 +0000 Subject: [PATCH] feat: add Bedrock Knowledge Base capability 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). --- pydantic_ai_harness/bedrock_kb/README.md | 95 ++++++ pydantic_ai_harness/bedrock_kb/__init__.py | 6 + pydantic_ai_harness/bedrock_kb/_capability.py | 84 ++++++ pydantic_ai_harness/bedrock_kb/_toolset.py | 270 ++++++++++++++++++ tests/bedrock_kb/__init__.py | 0 tests/bedrock_kb/test_bedrock_kb.py | 157 ++++++++++ 6 files changed, 612 insertions(+) create mode 100644 pydantic_ai_harness/bedrock_kb/README.md create mode 100644 pydantic_ai_harness/bedrock_kb/__init__.py create mode 100644 pydantic_ai_harness/bedrock_kb/_capability.py create mode 100644 pydantic_ai_harness/bedrock_kb/_toolset.py create mode 100644 tests/bedrock_kb/__init__.py create mode 100644 tests/bedrock_kb/test_bedrock_kb.py diff --git a/pydantic_ai_harness/bedrock_kb/README.md b/pydantic_ai_harness/bedrock_kb/README.md new file mode 100644 index 00000000..96b7d9e8 --- /dev/null +++ b/pydantic_ai_harness/bedrock_kb/README.md @@ -0,0 +1,95 @@ +# Bedrock Knowledge Base + +Connect your Pydantic AI agent to [Amazon Bedrock Managed Knowledge Bases](https://docs.aws.amazon.com/bedrock/latest/userguide/knowledge-base.html) for retrieval-augmented generation (RAG). + +## Features + +- **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) + +## Quick start + +```bash +uv add "pydantic-ai-harness[bedrock-kb]" +``` + +```python +from pydantic_ai import Agent +from pydantic_ai_harness.bedrock_kb import BedrockKnowledgeBase + +agent = Agent( + 'anthropic:claude-sonnet-4-20250514', + capabilities=[ + BedrockKnowledgeBase( + knowledge_base_id='YOUR_KB_ID', + region_name='us-west-2', + ), + ], +) + +result = agent.run_sync('What are the company policies on remote work?') +print(result.output) +``` + +## Configuration + +| Parameter | Description | Default | +|---|---|---| +| `knowledge_base_id` | Bedrock KB ID | Env: `KNOWLEDGE_BASE_ID` | +| `region_name` | AWS region | Env: `AWS_REGION` or `us-east-1` | +| `use_agentic_retrieval` | Use agentic multi-step retrieval | `True` | +| `number_of_results` | Max results per search | `5` | +| `data_source_id` | Data source ID (for ingestion) | Env: `BEDROCK_DATA_SOURCE_ID` | +| `data_source_type` | `"S3"` or `"CUSTOM"` | `"S3"` | +| `data_source_bucket` | S3 bucket (for S3 mode) | Env: `BEDROCK_DATA_SOURCE_BUCKET` | +| `include_ingest_tool` | Expose `ingest_document` tool | `False` | + +## Tools exposed to the agent + +### `search_knowledge_base(query: str) -> str` +Searches the KB for documents relevant to the query. Uses agentic retrieval by default. + +### `ingest_document(content, document_id, mime_type, s3_uri) -> str` +*(Only when `include_ingest_tool=True`)* + +Ingests a document into the KB. Three modes: +- **Inline text**: `content="Your text here"` +- **S3 reference**: `s3_uri="s3://bucket/file.pdf"` +- **Binary**: `content="", mime_type="application/pdf"` + +## IAM Permissions + +```json +{ + "Effect": "Allow", + "Action": [ + "bedrock:Retrieve", + "bedrock:AgenticRetrieveStream" + ], + "Resource": "arn:aws:bedrock:REGION:ACCOUNT:knowledge-base/KB_ID" +} +``` + +For ingestion, also add: +```json +{ + "Effect": "Allow", + "Action": "bedrock:IngestKnowledgeBaseDocuments", + "Resource": "arn:aws:bedrock:REGION:ACCOUNT:knowledge-base/KB_ID" +} +``` + +## Prerequisites + +- AWS credentials configured (via environment, IAM role, or AWS profile) +- A Bedrock Managed Knowledge Base created (via console, CDK, or CloudFormation) +- `boto3 >= 1.43.2` installed + +## References + +- [Ingest documents directly into a knowledge base](https://docs.aws.amazon.com/bedrock/latest/userguide/kb-direct-ingestion.html) +- [IngestKnowledgeBaseDocuments API](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent_IngestKnowledgeBaseDocuments.html) +- [Connect to a custom data source](https://docs.aws.amazon.com/bedrock/latest/userguide/custom-data-source-connector.html) +- [AgenticRetrieveStream API](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_AgenticRetrieveStream.html) diff --git a/pydantic_ai_harness/bedrock_kb/__init__.py b/pydantic_ai_harness/bedrock_kb/__init__.py new file mode 100644 index 00000000..80757e9c --- /dev/null +++ b/pydantic_ai_harness/bedrock_kb/__init__.py @@ -0,0 +1,6 @@ +"""Bedrock Knowledge Base capability: retrieval-augmented generation using Amazon Bedrock Managed Knowledge Bases.""" + +from pydantic_ai_harness.bedrock_kb._capability import BedrockKnowledgeBase +from pydantic_ai_harness.bedrock_kb._toolset import BedrockKBToolset + +__all__ = ['BedrockKnowledgeBase', 'BedrockKBToolset'] diff --git a/pydantic_ai_harness/bedrock_kb/_capability.py b/pydantic_ai_harness/bedrock_kb/_capability.py new file mode 100644 index 00000000..af3fc5c3 --- /dev/null +++ b/pydantic_ai_harness/bedrock_kb/_capability.py @@ -0,0 +1,84 @@ +"""Bedrock Knowledge Base capability: connect Pydantic AI agents to Amazon Bedrock Managed KBs.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from pydantic_ai.capabilities import AbstractCapability +from pydantic_ai.tools import AgentDepsT + +from pydantic_ai_harness.bedrock_kb._toolset import BedrockKBToolset + + +@dataclass +class BedrockKnowledgeBase(AbstractCapability[AgentDepsT]): + """Amazon Bedrock Knowledge Base retrieval capability. + + Gives agents access to a Bedrock Managed Knowledge Base for RAG. Supports: + - Agentic retrieval (multi-step, reasoning-enhanced search) + - Standard semantic retrieval + - Direct document ingestion (CUSTOM data source) + + Usage: + ```python + from pydantic_ai import Agent + from pydantic_ai_harness.bedrock_kb import BedrockKnowledgeBase + + agent = Agent( + 'anthropic:claude-sonnet-4-20250514', + capabilities=[ + BedrockKnowledgeBase( + knowledge_base_id='YOUR_KB_ID', + region_name='us-west-2', + ), + ], + ) + ``` + + References: + - https://docs.aws.amazon.com/bedrock/latest/userguide/knowledge-base.html + - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_AgenticRetrieveStream.html + - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent_IngestKnowledgeBaseDocuments.html + """ + + knowledge_base_id: str = '' + """Knowledge Base ID. Falls back to KNOWLEDGE_BASE_ID env var.""" + + region_name: str = '' + """AWS region. Falls back to AWS_REGION env var, then us-east-1.""" + + data_source_id: str = '' + """Data source ID for ingestion. Falls back to BEDROCK_DATA_SOURCE_ID env var.""" + + data_source_type: str = 'S3' + """Data source type: 'S3' (upload + sync) or 'CUSTOM' (direct ingestion via DLA).""" + + data_source_bucket: str = '' + """S3 bucket for S3 data source mode. Falls back to BEDROCK_DATA_SOURCE_BUCKET env var.""" + + use_agentic_retrieval: bool = True + """Use agentic retrieval (multi-step reasoning) instead of standard Retrieve.""" + + number_of_results: int = 5 + """Maximum number of results to return from retrieval.""" + + include_ingest_tool: bool = False + """Whether to expose the ingest_document tool to the agent.""" + + 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}') + + def get_toolset(self) -> BedrockKBToolset[AgentDepsT]: + """Build and return the Bedrock KB toolset.""" + return BedrockKBToolset[AgentDepsT]( + knowledge_base_id=self.knowledge_base_id, + region_name=self.region_name, + data_source_id=self.data_source_id, + data_source_type=self.data_source_type, + data_source_bucket=self.data_source_bucket, + use_agentic_retrieval=self.use_agentic_retrieval, + number_of_results=self.number_of_results, + include_ingest_tool=self.include_ingest_tool, + ) diff --git a/pydantic_ai_harness/bedrock_kb/_toolset.py b/pydantic_ai_harness/bedrock_kb/_toolset.py new file mode 100644 index 00000000..b9115929 --- /dev/null +++ b/pydantic_ai_harness/bedrock_kb/_toolset.py @@ -0,0 +1,270 @@ +"""Bedrock Knowledge Base toolset: retrieval and ingestion tools.""" + +from __future__ import annotations + +import base64 +import logging +import os +from typing import Any + +from pydantic_ai.exceptions import ModelRetry +from pydantic_ai.tools import AgentDepsT +from pydantic_ai.toolsets import FunctionToolset + +logger = logging.getLogger(__name__) + + +class BedrockKBToolset(FunctionToolset[AgentDepsT]): + """Toolset providing Amazon Bedrock Knowledge Base retrieval and ingestion. + + Tools: + - search_knowledge_base: Retrieve relevant documents for a query + - ingest_document: Ingest a document into the knowledge base (optional) + """ + + def __init__( + self, + *, + knowledge_base_id: str, + region_name: str, + data_source_id: str, + data_source_type: str, + data_source_bucket: str, + use_agentic_retrieval: bool, + number_of_results: int, + include_ingest_tool: bool, + ) -> None: + super().__init__() + self._knowledge_base_id = knowledge_base_id or os.environ.get('KNOWLEDGE_BASE_ID', '') + self._region_name = region_name or os.environ.get('AWS_REGION', 'us-east-1') + self._data_source_id = data_source_id or os.environ.get('BEDROCK_DATA_SOURCE_ID', '') + self._data_source_type = data_source_type.upper() + self._data_source_bucket = data_source_bucket or os.environ.get('BEDROCK_DATA_SOURCE_BUCKET', '') + self._use_agentic_retrieval = use_agentic_retrieval + self._number_of_results = number_of_results + self._runtime_client: Any = None + self._agent_client: Any = None + + self.add_function(self.search_knowledge_base, name='search_knowledge_base') + if include_ingest_tool: + self.add_function(self.ingest_document, name='ingest_document') + + @property + def _get_runtime_client(self) -> Any: + if self._runtime_client is None: + import boto3 + from botocore.config import Config + + self._runtime_client = boto3.client( + 'bedrock-agent-runtime', + region_name=self._region_name, + config=Config(user_agent_extra='pydantic-ai-harness/bedrock-kb'), + ) + return self._runtime_client + + @property + def _get_agent_client(self) -> Any: + if self._agent_client is None: + import boto3 + from botocore.config import Config + + self._agent_client = boto3.client( + 'bedrock-agent', + region_name=self._region_name, + config=Config(user_agent_extra='pydantic-ai-harness/bedrock-kb'), + ) + return self._agent_client + + async def search_knowledge_base(self, query: str) -> str: + """Search the knowledge base for documents relevant to the query. + + Args: + query: The natural language question or search query. + + Returns: + Formatted search results with content and metadata. + """ + if not self._knowledge_base_id: + raise ModelRetry('Knowledge Base ID is not configured. Set KNOWLEDGE_BASE_ID env var.') + + try: + if self._use_agentic_retrieval: + return await self._agentic_retrieve(query) + else: + return await self._standard_retrieve(query) + except Exception as e: + logger.warning('KB search failed: %s', e) + raise ModelRetry(f'Knowledge Base search failed: {e}') from e + + 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) + + 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 ingest_document( + self, + content: str, + document_id: str = '', + mime_type: str = '', + s3_uri: str = '', + ) -> str: + """Ingest a document into the knowledge base. + + Supports three modes: + - Inline text: provide `content` with plain text + - S3 reference: provide `s3_uri` to ingest an existing S3 object + - Binary: provide base64-encoded `content` with `mime_type` + + Args: + content: Document text content, or base64-encoded bytes if mime_type is set. + document_id: Optional identifier for the document. + mime_type: MIME type for binary content (e.g., 'application/pdf'). + s3_uri: S3 URI to ingest (e.g., 's3://bucket/key'). Overrides content. + + Returns: + Ingestion status message. + """ + if not self._knowledge_base_id: + raise ModelRetry('Knowledge Base ID is not configured.') + if not self._data_source_id: + raise ModelRetry('Data source ID is not configured for ingestion.') + + import asyncio + import uuid + + doc_id = document_id or str(uuid.uuid4()) + + def _call() -> str: + if self._data_source_type == 'CUSTOM': + return self._ingest_direct(doc_id, content, mime_type, s3_uri) + else: + return self._ingest_s3(doc_id, content, mime_type) + + loop = asyncio.get_event_loop() + try: + return await loop.run_in_executor(None, _call) + except Exception as e: + logger.warning('Ingestion failed: %s', e) + raise ModelRetry(f'Document ingestion failed: {e}') from e + + def _ingest_direct(self, doc_id: str, content: str, mime_type: str, s3_uri: str) -> str: + """Ingest via IngestKnowledgeBaseDocuments API (CUSTOM data source).""" + if s3_uri: + doc = { + 'content': { + 'dataSourceType': 'CUSTOM', + 'custom': { + 'customDocumentIdentifier': {'id': doc_id}, + 'sourceType': 'S3_LOCATION', + 's3Location': {'uri': s3_uri}, + }, + }, + } + elif mime_type: + doc = { + 'content': { + 'dataSourceType': 'CUSTOM', + 'custom': { + 'customDocumentIdentifier': {'id': doc_id}, + 'sourceType': 'IN_LINE', + 'inlineContent': { + 'type': 'BYTE', + 'byteContent': {'data': content, 'mimeType': mime_type}, + }, + }, + }, + } + else: + doc = { + 'content': { + 'dataSourceType': 'CUSTOM', + 'custom': { + 'customDocumentIdentifier': {'id': doc_id}, + 'sourceType': 'IN_LINE', + 'inlineContent': { + 'type': 'TEXT', + 'textContent': {'data': content}, + }, + }, + }, + } + + response = self._get_agent_client.ingest_knowledge_base_documents( + knowledgeBaseId=self._knowledge_base_id, + dataSourceId=self._data_source_id, + documents=[doc], + ) + status = response.get('documentDetails', [{}])[0].get('status', 'UNKNOWN') + return f'Document "{doc_id}" ingestion started. Status: {status}' + + def _ingest_s3(self, doc_id: str, content: str, mime_type: str) -> str: + """Ingest via S3 upload + StartIngestionJob.""" + import boto3 + + if not self._data_source_bucket: + return 'Error: No S3 bucket configured for ingestion.' + + s3 = boto3.client('s3', region_name=self._region_name) + key = f'pydantic-ai-harness/{doc_id}.txt' + s3.put_object( + Bucket=self._data_source_bucket, + Key=key, + Body=content.encode('utf-8'), + ContentType=mime_type or 'text/plain', + ) + + self._get_agent_client.start_ingestion_job( + knowledgeBaseId=self._knowledge_base_id, + dataSourceId=self._data_source_id, + description=f'Ingest {doc_id} via pydantic-ai-harness', + ) + return f'Document "{doc_id}" uploaded to s3://{self._data_source_bucket}/{key} and ingestion started.' diff --git a/tests/bedrock_kb/__init__.py b/tests/bedrock_kb/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/bedrock_kb/test_bedrock_kb.py b/tests/bedrock_kb/test_bedrock_kb.py new file mode 100644 index 00000000..235440ab --- /dev/null +++ b/tests/bedrock_kb/test_bedrock_kb.py @@ -0,0 +1,157 @@ +"""Tests for Bedrock Knowledge Base capability.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +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')) + + +class TestBedrockKnowledgeBaseCapability: + """Tests for the BedrockKnowledgeBase capability class.""" + + def test_default_values(self): + from pydantic_ai_harness.bedrock_kb._capability import BedrockKnowledgeBase + + kb = BedrockKnowledgeBase() + assert kb.use_agentic_retrieval is True + assert kb.number_of_results == 5 + assert kb.data_source_type == 'S3' + assert kb.include_ingest_tool is False + + def test_invalid_number_of_results(self): + from pydantic_ai_harness.bedrock_kb._capability import BedrockKnowledgeBase + + with pytest.raises(ValueError, match='number_of_results must be positive'): + BedrockKnowledgeBase(number_of_results=0) + + def test_get_toolset_returns_toolset(self): + from pydantic_ai_harness.bedrock_kb._capability import BedrockKnowledgeBase + + kb = BedrockKnowledgeBase(knowledge_base_id='TEST', region_name='us-west-2') + toolset = kb.get_toolset() + assert toolset._knowledge_base_id == 'TEST' + assert toolset._region_name == 'us-west-2'