Skip to main content

Bedrock Knowledge Bases

AWS Bedrock Knowledge Bases allows you to connect your LLM's to your organization's data, letting your models retrieve and reference information specific to your business.

PropertyDetails
DescriptionBedrock Knowledge Bases connects your data to LLM's, enabling them to retrieve and reference your organization's information in their responses.
Provider Route on LiteLLMbedrock in the litellm vector_store_registry
Provider DocAWS Bedrock Knowledge Bases ↗

Quick Start​

LiteLLM Python SDK​

Example using LiteLLM Python SDK
import os
import litellm

from litellm.vector_stores.vector_store_registry import VectorStoreRegistry, LiteLLM_ManagedVectorStore

# Init vector store registry with your Bedrock Knowledge Base
litellm.vector_store_registry = VectorStoreRegistry(
vector_stores=[
LiteLLM_ManagedVectorStore(
vector_store_id="YOUR_KNOWLEDGE_BASE_ID", # KB ID from AWS Bedrock
custom_llm_provider="bedrock"
)
]
)

# Make a completion request using your Knowledge Base
response = await litellm.acompletion(
model="anthropic/claude-3-5-sonnet",
messages=[{"role": "user", "content": "What does our company policy say about remote work?"}],
tools=[
{
"type": "file_search",
"vector_store_ids": ["YOUR_KNOWLEDGE_BASE_ID"]
}
],
)

print(response.choices[0].message.content)

LiteLLM Proxy​

1. Configure your vector_store_registry​

model_list:
- model_name: claude-3-5-sonnet
litellm_params:
model: anthropic/claude-3-5-sonnet
api_key: os.environ/ANTHROPIC_API_KEY

vector_store_registry:
- vector_store_name: "bedrock-company-docs"
litellm_params:
vector_store_id: "YOUR_KNOWLEDGE_BASE_ID"
custom_llm_provider: "bedrock"
vector_store_description: "Bedrock Knowledge Base for company documents"
vector_store_metadata:
source: "Company internal documentation"

2. Make a request with vector_store_ids parameter​

curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer $LITELLM_API_KEY" \
-d '{
"model": "claude-3-5-sonnet",
"messages": [{"role": "user", "content": "What does our company policy say about remote work?"}],
"tools": [
{
"type": "file_search",
"vector_store_ids": ["YOUR_KNOWLEDGE_BASE_ID"]
}
]
}'

Futher Reading Vector Stores: