mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix(llms): resolve filename/file_id from location field when metadata URI absent
Bedrock Knowledge Base results backed by S3 (and other sources) sometimes omit `x-amz-bedrock-kb-source-uri` from the metadata field but always populate the `location` field. Without this fix, `filename` and `file_id` fall back to `bedrock-kb-document-<data-source-id>` instead of the actual file path.
This commit is contained in:
Hai Huynhlam
parent
ca4d4a0188
commit
f29972f3b7
2 changed files with 190 additions and 4 deletions
|
|
@ -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": <key>, "value": <value>, "operator": <operator>} OR {"and" | "or": [{"key": <key>, "value": <value>, "operator": <operator>}]}
|
||||
# aws kb filters = {"operator": {"<key>": <value>}} OR {"andAll | orAll": [{"operator": {"<key>": <value>}}]}
|
||||
|
|
@ -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-<id> 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(
|
||||
|
|
|
|||
|
|
@ -24,4 +24,151 @@ def test_transform_search_request():
|
|||
)
|
||||
|
||||
assert url.endswith("/kb123/retrieve")
|
||||
assert body["retrievalQuery"].get("text") == "hello"
|
||||
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"
|
||||
Loading…
Add table
Reference in a new issue