diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 72e1e1470d3..235399f7396 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -152,7 +152,6 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): if param == "max_num_results": optional_params["numberOfResults"] = value elif param == "filters" and value is not None: - # map the openai filters to the aws kb filters format # openai filters = {"key": , "value": , "operator": } OR {"and" | "or": [{"key": , "value": , "operator": }]} # aws kb filters = {"operator": {"": }} OR {"andAll | orAll": [{"operator": {"": }}]} @@ -297,6 +296,34 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): ) return f"bedrock-kb-document-{data_source_id}" + def _get_uri_from_location(self, location: Dict[str, Any]) -> Optional[str]: + """ + Extract source URI from Bedrock KB location field. + + Supports all location types from the Bedrock Retrieve API: + https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_Retrieve.html + """ + if not location: + return None + location_type = (location.get("type") or "").upper() + type_map = { + "CONFLUENCE": ("confluenceLocation", "url"), + "CUSTOM": ("customDocumentLocation", "id"), + "KENDRA": ("kendraDocumentLocation", "uri"), + "S3": ("s3Location", "uri"), + "SALESFORCE": ("salesforceLocation", "url"), + "SHAREPOINT": ("sharePointLocation", "url"), + "WEB": ("webLocation", "url"), + # SQL sources expose only a query string, not a stable document URI; + # callers receive the bedrock-kb-document- fallback instead. + } + entry = type_map.get(location_type) + if not entry: + return None + loc_key, uri_key = entry + loc_data = location.get(loc_key) or {} + return loc_data.get(uri_key) or None + def _get_attributes_from_metadata(self, metadata: Dict[str, Any]) -> Dict[str, Any]: """ Extract all attributes from Bedrock KB metadata. @@ -320,8 +347,20 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): # Extract metadata and use helper functions metadata = item.get("metadata", {}) or {} - file_id = self._get_file_id_from_metadata(metadata) - filename = self._get_filename_from_metadata(metadata) + # Resolve source URI from location field if not present in metadata. + # Use a separate dict for file_id/filename resolution so the + # synthesized key does not leak into attributes. + source_uri = metadata.get("x-amz-bedrock-kb-source-uri") + if not source_uri: + location = item.get("location", {}) or {} + source_uri = self._get_uri_from_location(location) + metadata_for_id = ( + {**metadata, "x-amz-bedrock-kb-source-uri": source_uri} + if source_uri + else metadata + ) + file_id = self._get_file_id_from_metadata(metadata_for_id) + filename = self._get_filename_from_metadata(metadata_for_id) attributes = self._get_attributes_from_metadata(metadata) results.append( diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index 28b60e5e75f..fa2d5693c16 100644 --- a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -24,4 +24,151 @@ def test_transform_search_request(): ) assert url.endswith("/kb123/retrieve") - assert body["retrievalQuery"].get("text") == "hello" \ No newline at end of file + assert body["retrievalQuery"].get("text") == "hello" + + +def test_get_uri_from_location_s3(): + config = BedrockVectorStoreConfig() + location = { + "type": "S3", + "s3Location": {"uri": "s3://my-bucket/docs/file.pdf"}, + } + assert config._get_uri_from_location(location) == "s3://my-bucket/docs/file.pdf" + + +def test_get_uri_from_location_web(): + config = BedrockVectorStoreConfig() + location = { + "type": "WEB", + "webLocation": {"url": "https://example.com/page"}, + } + assert config._get_uri_from_location(location) == "https://example.com/page" + + +def test_get_uri_from_location_confluence(): + config = BedrockVectorStoreConfig() + location = { + "type": "CONFLUENCE", + "confluenceLocation": {"url": "https://myorg.atlassian.net/wiki/spaces/PROJ/pages/123"}, + } + assert config._get_uri_from_location(location) == "https://myorg.atlassian.net/wiki/spaces/PROJ/pages/123" + + +def test_get_uri_from_location_kendra(): + config = BedrockVectorStoreConfig() + location = { + "type": "KENDRA", + "kendraDocumentLocation": {"uri": "kendra://index-id/doc-id"}, + } + assert config._get_uri_from_location(location) == "kendra://index-id/doc-id" + + +def test_get_uri_from_location_salesforce(): + config = BedrockVectorStoreConfig() + location = { + "type": "SALESFORCE", + "salesforceLocation": {"url": "https://myorg.salesforce.com/articles/example"}, + } + assert config._get_uri_from_location(location) == "https://myorg.salesforce.com/articles/example" + + +def test_get_uri_from_location_sharepoint(): + config = BedrockVectorStoreConfig() + location = { + "type": "SHAREPOINT", + "sharePointLocation": {"url": "https://myorg.sharepoint.com/sites/team/doc.docx"}, + } + assert config._get_uri_from_location(location) == "https://myorg.sharepoint.com/sites/team/doc.docx" + + +def test_get_uri_from_location_custom(): + config = BedrockVectorStoreConfig() + location = { + "type": "CUSTOM", + "customDocumentLocation": {"id": "custom-doc-id-abc123"}, + } + assert config._get_uri_from_location(location) == "custom-doc-id-abc123" + + +def test_get_uri_from_location_unknown_returns_none(): + config = BedrockVectorStoreConfig() + assert config._get_uri_from_location({}) is None + assert config._get_uri_from_location({"type": "UNKNOWN"}) is None + assert config._get_uri_from_location({"type": "S3"}) is None # missing s3Location + + +def test_transform_response_uses_location_uri(): + """ + When x-amz-bedrock-kb-source-uri is absent from metadata, the URI should + be resolved from location.s3Location.uri and used for filename/file_id. + """ + config = BedrockVectorStoreConfig() + mock_log = MagicMock() + mock_log.model_call_details = {"query": "test query"} + + raw_response = { + "retrievalResults": [ + { + "content": {"text": "some content", "type": "TEXT"}, + "location": { + "s3Location": {"uri": "s3://my-company-bedrock-kb/docs/document.md"}, + "type": "S3", + }, + "metadata": { + "x-amz-bedrock-kb-source-file-modality": "TEXT", + "x-amz-bedrock-kb-chunk-id": "8befd6d7-d8d1-49f6-b01d-54cf4e77c01a", + "x-amz-bedrock-kb-data-source-id": "ABCDE12345", + }, + "score": 0.506902021031646, + } + ] + } + + mock_http_response = MagicMock() + mock_http_response.json.return_value = raw_response + mock_http_response.status_code = 200 + + result = config.transform_search_vector_store_response(mock_http_response, mock_log) + + assert len(result["data"]) == 1 + item = result["data"][0] + assert item["file_id"] == "s3://my-company-bedrock-kb/docs/document.md" + assert item["filename"] == "document.md" + # Synthesized URI must not leak into attributes + assert "x-amz-bedrock-kb-source-uri" not in item["attributes"] + + +def test_transform_response_metadata_uri_takes_precedence(): + """ + When x-amz-bedrock-kb-source-uri is already in metadata, it must be used + and the location field must be ignored. + """ + config = BedrockVectorStoreConfig() + mock_log = MagicMock() + mock_log.model_call_details = {"query": "test query"} + + raw_response = { + "retrievalResults": [ + { + "content": {"text": "some content", "type": "TEXT"}, + "location": { + "s3Location": {"uri": "s3://example-bucket/location-path/file.pdf"}, + "type": "S3", + }, + "metadata": { + "x-amz-bedrock-kb-source-uri": "s3://example-bucket/metadata-path/other.pdf", + }, + "score": 0.9, + } + ] + } + + mock_http_response = MagicMock() + mock_http_response.json.return_value = raw_response + mock_http_response.status_code = 200 + + result = config.transform_search_vector_store_response(mock_http_response, mock_log) + + item = result["data"][0] + assert item["file_id"] == "s3://example-bucket/metadata-path/other.pdf" + assert item["filename"] == "other.pdf" \ No newline at end of file