Skip to content

feat(retrieve): expand document location extraction for all KB source types#366

Open
strands-agent wants to merge 1 commit intostrands-agents:mainfrom
strands-agent:feat/expand-retrieve-location-types
Open

feat(retrieve): expand document location extraction for all KB source types#366
strands-agent wants to merge 1 commit intostrands-agents:mainfrom
strands-agent:feat/expand-retrieve-location-types

Conversation

@strands-agent
Copy link
Contributor

Description

Expand the format_results_for_display function in the retrieve tool to handle all RetrievalResultLocation types from the Bedrock Knowledge Base API, not just customDocumentLocation and s3Location.

Changes:

  • Added support for webLocation (WEB)
  • Added support for confluenceLocation (CONFLUENCE)
  • Added support for salesforceLocation (SALESFORCE)
  • Added support for sharePointLocation (SHAREPOINT)
  • Added support for kendraDocumentLocation (KENDRA)
  • Added support for sqlLocation (SQL)
  • Added 7 new unit tests covering all location types plus unknown/future types

Reference:

https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_RetrievalResultLocation.html

Related Issues

Closes #364

Type of Change

  • Bug fix
  • New Feature (adds support for additional location types)
  • Breaking change
  • Other

Testing

  • All 30 tests pass including 7 new tests
  • Ran ruff format and ruff check - no issues
$ python -m pytest tests/test_retrieve.py -v
============================= test session starts ==============================
collected 30 items

tests/test_retrieve.py::test_format_results_with_web_location PASSED     [ 80%]
tests/test_retrieve.py::test_format_results_with_confluence_location PASSED [ 83%]
tests/test_retrieve.py::test_format_results_with_salesforce_location PASSED [ 86%]
tests/test_retrieve.py::test_format_results_with_sharepoint_location PASSED [ 90%]
tests/test_retrieve.py::test_format_results_with_kendra_location PASSED  [ 93%]
tests/test_retrieve.py::test_format_results_with_sql_location PASSED     [ 96%]
tests/test_retrieve.py::test_format_results_with_unknown_location PASSED [100%]

============================== 30 passed in 0.79s ==============================

Checklist

  • I have read the CONTRIBUTING document
  • I have added tests that prove my fix is effective or my feature works
  • I have updated the documentation accordingly (code comments reference AWS API docs)
  • My changes generate no new warnings
  • Any dependent changes have been merged and published

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

… types

Add support for all RetrievalResultLocation types from the Bedrock API:
- webLocation (WEB)
- confluenceLocation (CONFLUENCE)
- salesforceLocation (SALESFORCE)
- sharePointLocation (SHAREPOINT)
- kendraDocumentLocation (KENDRA)
- sqlLocation (SQL)

Previously only customDocumentLocation and s3Location were handled,
which limited usability for knowledge bases built on Kendra or other
data sources like Confluence, SharePoint, etc.

Added tests for each new location type plus a test for unknown/future
location types to ensure graceful handling.

Closes strands-agents#364
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[FEATURE] Expand document location extraction for retrieve tool

1 participant