AWS Builder Center

Builder Center

Building hybrid retrieval for AI agents with Bedrock and DynamoDB

Combine vector search with structured entity queries across two data stores, letting an AI agent decide which retrieval path fits each query.

Kajal

Vector search is great at answering questions like "what are some healthy dinner ideas?". It finds semantically relevant passages even when the exact wording differs between the query and the source material. But if you were to ask "which recipes use chicken and take less than 30 minutes?", vector search struggles. That second question requires structured filtering across known attributes and their relationships. This is the exact kind of query a key-value store handles in milliseconds.
Classically, retrieval-augmented generation (RAG) applications start with vector-only search and hit this wall within weeks of production use. Your users ask questions that blend semantic understanding with entity-specific lookups, and a single retrieval mechanism can't serve both patterns well. The common response is to add metadata filtering to the vector store, however this approach can break down when entity relationships become complex or when you need precise boolean logic across multiple dimensions.
In this post, we demonstrate a hybrid retrieval pattern that pairs Amazon Bedrock Knowledge Bases for semantic search with Amazon DynamoDB for structured entity queries. An AI agent built with the Strands Agents SDK decides which retrieval path to use (or combines both) based on the nature of each query. This differs from "hybrid search" (which combines semantic and keyword search within a single store). Hybrid retrieval routes queries across two purpose-built stores, each optimized for its access pattern.

Solution overview

The architecture splits retrieval into two complementary paths. Amazon Bedrock Knowledge Bases indexes full-text documents and serves semantic similarity queries through vector embeddings. Amazon DynamoDB stores structured entity data. These include documents, topics, technologies, and categories. DynamoDB additionally includes access patterns designed for precise lookups and relationship traversal.
An AI agent deployed on Amazon Bedrock AgentCore Runtime sits between the user and both data stores. The agent's system prompt teaches it to recognize query types and pick the right retrieval tool. For complex questions that need both semantic context and entity precision, the agent calls both tools and merges the results into a unified response.
The agent exposes three tools to do this: search_knowledge_base for semantic retrieval, query_entities for structured lookups, and compare_findings for cross-source synthesis. The foundation model powering the agent figures out which tool (or combination) best serves each query.
Architecture diagram showing a user query flowing to an AI agent on Amazon Bedrock AgentCore Runtime with three tools: search_knowledge_base connects to Amazon Bedrock Knowledge Base for vector search, query_entities connects to Amazon DynamoDB for structured queries, and compare_findings calls both stores and merges results.
The AI agent routes each query to the appropriate retrieval path based on whether the question requires semantic search, structured entity lookup, or both.

Prerequisites

To implement this pattern, you need the following:
  • An AWS account with access to Amazon Bedrock (foundation models and Knowledge Bases enabled)
  • An Amazon Bedrock knowledge base configured with your document collection
  • An Amazon DynamoDB table with entity data (schema described in the next section)
  • Python 3.12 or later
  • The Strands Agents SDK installed (pip install strands-agents strands-agents-tools)
  • The AWS SDK for Python (pip install boto3)
  • IAM permissions for bedrock-agent-runtime:Retrieve, dynamodb:Query, and bedrock:InvokeModel

The DynamoDB single-table design

A single DynamoDB table stores all entity data using a composite key structure. The partition key (PK) and sort key (SK) encode entity type and relationship information, so one table handles documents, entities, and the connections between them without needing additional tables.
The table stores three item types:
  • Document metadata stores bibliographic information. A research paper might be stored as PK: DOC#paper-001, SK: METADATA.
  • Entity mentions capture which entities appear in which documents. If that paper mentions the brain region "hippocampus," the item is stored as PK: DOC#paper-001, SK: MENTIONS#BRAIN_REGION#hippocampus.
  • Entity metadata describes the entities themselves. The hippocampus entity is stored as PK: BRAIN_REGION#hippocampus, SK: METADATA.
This structure gives you one access pattern directly: query PK: DOC#paper-001 with a SK prefix of MENTIONS# and you get every entity mentioned in that paper.
A Global Secondary Index (GSI1) enables two more access patterns that the main table keys cannot support:
  • List all entities in a category. For example, "show me every brain region in the system." The GSI uses GSI1PK: CATEGORY#BRAIN_REGION with GSI1SK: BRAIN_REGION#hippocampus, BRAIN_REGION#amygdala, and so on.
  • Reverse lookup from entity to documents. For example, "which papers mention hippocampus?" The GSI uses GSI1PK: BRAIN_REGION#hippocampus with GSI1SK: DOC#paper-001, DOC#paper-047, and so on. This is the key pattern that vector search alone cannot provide, because it requires traversing a known relationship rather than finding semantic similarity.
In practice, the query_entities tool combines all three: listing entities in a category via GSI1, finding all mentions within a specific document via the main table, and reverse-looking up all documents for a specific entity via GSI1.
Here is how we write these items using DynamoDB's batch_writer:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
@tool
def store_entities(doc_id: str, title: str, entities: str) -> str:
"""Store extracted entities in DynamoDB with document references."""
entity_data = json.loads(entities)
items_written = 0

with table.batch_writer() as batch:
# Document metadata item
batch.put_item(Item={
"PK": f"DOC#{doc_id}",
"SK": "METADATA",
"title": title,
"doc_id": doc_id,
"entity_count": sum(len(v) for v in entity_data.values()),
"GSI1PK": "DOCUMENTS",
"GSI1SK": f"DOC#{doc_id}",
})
items_written += 1

for category, entity_list in entity_data.items():
for entity_name in entity_list:
entity_id = entity_name.lower().strip().replace(" ", "-")

# Entity metadata (idempotent, overwrites if exists)
batch.put_item(Item={
"PK": f"{category.upper()}#{entity_id}",
"SK": "METADATA",
"entity_name": entity_name,
"category": category,
"GSI1PK": f"CATEGORY#{category}",
"GSI1SK": f"{category.upper()}#{entity_id}",
})
items_written += 1

# Adjacency list: document → entity relationship
batch.put_item(Item={
"PK": f"DOC#{doc_id}",
"SK": f"MENTIONS#{category.upper()}#{entity_id}",
"entity_name": entity_name,
"category": category,
"GSI1PK": f"{category.upper()}#{entity_id}",
"GSI1SK": f"DOC#{doc_id}",
})
items_written += 1

return json.dumps({"status": "success", "items_written": items_written})

Building the retrieval tools

Each tool is a Python function decorated with @tool from the Strands Agents SDK. The decorator registers the function so the agent can call it during conversation. Let's walk through each one.
The search_knowledge_base tool wraps the Amazon Bedrock Agent Runtime Retrieve API. It takes a natural language query, runs vector similarity search against the knowledge base, and returns the top results with source metadata. The reason this tool works best on exploratory questions is its tolerance for vocabulary mismatch. A query about "data caching strategies" still finds documents about "ElastiCache best practices."
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
@tool
def search_knowledge_base(query: str, max_results: int = 5) -> str:
"""Search the document knowledge base using semantic search."""
response = bedrock_agent_runtime.retrieve(
knowledgeBaseId=KNOWLEDGE_BASE_ID,
retrievalQuery={"text": query},
retrievalConfiguration={
"vectorSearchConfiguration": {
"numberOfResults": max_results,
},
},
)

results = []
for item in response.get("retrievalResults", []):
content = item.get("content", {}).get("text", "")
s3_location = item.get("location", {}).get("s3Location", {})
results.append({
"text": content,
"source": s3_location.get("uri", ""),
"score": item.get("score", 0),
})

return json.dumps({"query": query, "results": results, "count": len(results)})
The query_entities tool gives us three structured access patterns against DynamoDB. You specify which pattern to use, and it constructs the right query. This is where DynamoDB comes in: each query returns exact matches in single-digit milliseconds, no matter how large the dataset gets.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
@tool
def query_entities(
query_type: str,
entity_category: str = "",
entity_name: str = "",
doc_id: str = "",
) -> str:
"""Query structured entities from the data store."""
if query_type == "by_category" and entity_category:
# List all entities in a category via GSI1
response = table.query(
IndexName="GSI1",
KeyConditionExpression=Key("GSI1PK").eq(f"CATEGORY#{entity_category}"),
)
entities = [
{"name": item.get("entity_name", ""), "id": item.get("SK", "")}
for item in response.get("Items", [])
if item.get("SK") == "METADATA"
]
return json.dumps({"category": entity_category, "entities": entities})

elif query_type == "by_document" and doc_id:
# List all entities extracted from a specific document
response = table.query(
KeyConditionExpression=(
Key("PK").eq(f"DOC#{doc_id}")
& Key("SK").begins_with("MENTIONS#")
),
)
entities = [
{"name": item.get("entity_name", ""), "category": item.get("category", "")}
for item in response.get("Items", [])
]
return json.dumps({"doc_id": doc_id, "entities": entities})

elif query_type == "by_entity" and entity_category and entity_name:
# Reverse lookup: find all documents mentioning a specific entity
entity_id = entity_name.lower().strip().replace(" ", "-")
response = table.query(
IndexName="GSI1",
KeyConditionExpression=Key("GSI1PK").eq(
f"{entity_category.upper()}#{entity_id}"
),
)
docs = [
item.get("GSI1SK", "").replace("DOC#", "")
for item in response.get("Items", [])
if item.get("GSI1SK", "").startswith("DOC#")
]
return json.dumps({"entity": entity_name, "documents": docs})
The compare_findings tool is where the hybrid pattern comes together. It runs a vector search through the knowledge base and a structured query through DynamoDB in the same call, then hands both result sets back to the agent for synthesis. Use this when a question needs both semantic context and entity precision, like "how do the documents that mention Lambda discuss cold start optimization?"
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
@tool
def compare_findings(topic: str, entity_category: str = "") -> str:
"""Compare findings across documents using both retrieval methods."""
comparison = {"topic": topic, "kb_results": [], "entity_data": []}

# Semantic search via Knowledge Base
response = bedrock_agent_runtime.retrieve(
knowledgeBaseId=KNOWLEDGE_BASE_ID,
retrievalQuery={"text": topic},
retrievalConfiguration={
"vectorSearchConfiguration": {"numberOfResults": 10},
},
)
for item in response.get("retrievalResults", []):
comparison["kb_results"].append({
"text": item.get("content", {}).get("text", ""),
"source": item.get("location", {}).get("s3Location", {}).get("uri", ""),
"score": item.get("score", 0),
})

# Structured entity data via DynamoDB
if entity_category:
response = table.query(
IndexName="GSI1",
KeyConditionExpression=Key("GSI1PK").eq(f"CATEGORY#{entity_category}"),
)
for item in response.get("Items", []):
if item.get("SK") == "METADATA":
comparison["entity_data"].append({
"name": item.get("entity_name", ""),
"category": entity_category,
})

return json.dumps(comparison)

Wiring the agent

The system prompt is what makes or breaks this pattern. Without explicit routing guidance, the model defaults to calling search_knowledge_base for everything, and that defeats the whole purpose of having two retrieval paths. We need to teach the agent when to use which tool.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
SYSTEM_PROMPT = """You are a research assistant with access to a document
library through two retrieval methods:

1. Vector search (search_knowledge_base) for semantic queries
2. Structured entity queries (query_entities) for entity lookups

Routing rules:
- If the user names specific entities (services, technologies, error types),
use query_entities first to find relevant documents, then
search_knowledge_base for context from those documents.
- If the question is exploratory or thematic, use search_knowledge_base.
- If the question involves comparing entities across documents,
use compare_findings.
- For ambiguous queries, start with search_knowledge_base and refine
with query_entities if the results reference specific entities.

Always cite source documents by title when referencing specific findings."""

agent = Agent(
model=BedrockModel(model_id="us.anthropic.claude-sonnet-4-20250514-v1:0"),
system_prompt=SYSTEM_PROMPT,
tools=[search_knowledge_base, query_entities, compare_findings],
)
The routing rules act as a decision tree for the model. It evaluates the user's query against these rules and picks the right tool. For ambiguous queries, we give it a fallback strategy: semantic search first, then structured refinement. That way we avoid empty results while still surfacing entity-specific data when it's relevant.

When to use which retrieval path

Here's a practical framework for deciding which path to take.
Vector search through Amazon Bedrock Knowledge Bases works best for exploratory questions where the user doesn't know exact terminology, for thematic queries that span multiple documents, and for questions about methods or concepts that may be described differently across sources. Vector search works best when vocabulary mismatch is likely. A query about "scaling web apps" still finds documents about "horizontal auto-scaling with Amazon EC2."
Structured queries through Amazon DynamoDB work best when the user names specific entities, when the question involves boolean logic across entity relationships (documents that mention technology A AND problem B), and when the answer requires traversing known relationships rather than discovering new ones. DynamoDB works best when you need exact matches fast, and it delivers in single-digit milliseconds regardless of dataset size.
The combined path through compare_findings serves questions that need both. We identify relevant documents through entity relationships (structured), then extract and synthesize thematic content from those documents (semantic). This works particularly well for review tasks where someone wants to understand how a specific topic is treated across a known set of documents.
A practical heuristic: if you can express the query as a DynamoDB key condition, use structured retrieval. If the query requires understanding meaning and context, use vector search. If it needs both, combine them.

Clean up

To avoid ongoing charges, delete the following resources if you deployed this pattern for testing:
  • The Amazon DynamoDB table
  • The Amazon Bedrock knowledge base and its associated Amazon S3 data source bucket
  • The Amazon Bedrock AgentCore Runtime agent (if deployed)
  • Any IAM roles created specifically for this pattern

Conclusion

Vector search and structured queries solve fundamentally different retrieval problems. If you try to force one mechanism to handle both, you end up with either imprecise entity lookups or poor semantic understanding. The hybrid retrieval pattern we walked through gives an AI agent the right tool for each query type, with a combined path for questions that span both domains.
This pattern scales naturally. Adding new entity categories to DynamoDB requires no changes to the vector store. Expanding the document corpus in Amazon Bedrock Knowledge Bases requires no schema changes in DynamoDB. The agent's system prompt is the only component that needs updating as you add new retrieval capabilities, and that's a one-line change, not a re-architecture.
To get started, configure an Amazon Bedrock knowledge base with your document corpus, model your entity relationships in a DynamoDB single-table design, and wire both into an agent using the Strands Agents SDK. The Amazon Bedrock Knowledge Bases documentation , Amazon DynamoDB single-table design guide , and Strands Agents SDK  provide the foundation for each component.
# amazon-bedrock# amazon-dynamodb# retreival-augmented-generation# ai-agents# strands-agents
Any opinions in this article are those of the individual author and may not reflect the opinions of AWS.
Enjoyed reading this content? Let the author know!

Your likes, comments, shares, and saves help creators reach more builders.

Loading recommendations

Loading article