(feat) Add Bedrock knowledge base pass through endpoints (#7267)
All checks were successful
Read Version from pyproject.toml / read-version (push) Successful in 56s

* bugfix: Proxy Routing for Bedrock Knowledgebase URLs are incorrect (#7097)

* Fixing routing bug where bedrock knowledgebase urls were being generated incorrectly

* Preparing for PR

* Preparing for PR

* Preparing for PR

---------

Co-authored-by: Luke Birk <lb0737@att.com>

* fix _is_bedrock_agent_runtime_route

* docs - Query Knowledge Base

* test_is_bedrock_agent_runtime_route

* fix bedrock_proxy_route

---------

Co-authored-by: LBirk <2731718+LBirk@users.noreply.github.com>
Co-authored-by: Luke Birk <lb0737@att.com>
This commit is contained in:
Ishaan Jaff 2024-12-16 22:19:34 -08:00 committed by GitHub
parent 3c984ed60e
commit f3b13a9af3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 51 additions and 2 deletions

View file

@ -32,6 +32,7 @@ from starlette.datastructures import QueryParams
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.batches.main import FileObject
from litellm.constants import BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES
from litellm.fine_tuning.main import vertex_fine_tuning_apis_instance
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -247,7 +248,7 @@ async def bedrock_proxy_route(
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
aws_region_name = litellm.utils.get_secret(secret_name="AWS_REGION_NAME")
if endpoint.startswith("agents/"): # handle bedrock agents
if _is_bedrock_agent_runtime_route(endpoint=endpoint): # handle bedrock agents
base_target_url = (
f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com"
)
@ -303,6 +304,16 @@ async def bedrock_proxy_route(
return received_value
def _is_bedrock_agent_runtime_route(endpoint: str) -> bool:
"""
Return True, if the endpoint should be routed to the `bedrock-agent-runtime` endpoint.
"""
for _route in BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES:
if _route in endpoint:
return True
return False
@router.api_route(
"/azure/{endpoint:path}",
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],