diff --git a/llama-index-integrations/readers/llama-index-readers-confluence/README.md b/llama-index-integrations/readers/llama-index-readers-confluence/README.md index 11ae8dbfed39be246409bfaa33ea8f8ef4ca4335..618780ea07e1be8a57a1d54c436f0c3847fcb962 100644 --- a/llama-index-integrations/readers/llama-index-readers-confluence/README.md +++ b/llama-index-integrations/readers/llama-index-readers-confluence/README.md @@ -68,7 +68,11 @@ base_url = "https://yoursite.atlassian.com/wiki" page_ids = ["<page_id_1>", "<page_id_2>", "<page_id_3"] space_key = "<space_key>" -reader = ConfluenceReader(base_url=base_url, oauth2=oauth2_dict) +reader = ConfluenceReader( + base_url=base_url, + oauth2=oauth2_dict, + client_args={"backoff_and_retry": True}, +) documents = reader.load_data( space_key=space_key, include_attachments=True, page_status="current" ) diff --git a/llama-index-integrations/readers/llama-index-readers-confluence/llama_index/readers/confluence/base.py b/llama-index-integrations/readers/llama-index-readers-confluence/llama_index/readers/confluence/base.py index b6c3278b9c1a516f015e98eace6eb5edca666b80..973da486ea703647f0ca0e437a291deb733926ea 100644 --- a/llama-index-integrations/readers/llama-index-readers-confluence/llama_index/readers/confluence/base.py +++ b/llama-index-integrations/readers/llama-index-readers-confluence/llama_index/readers/confluence/base.py @@ -33,6 +33,7 @@ class ConfluenceReader(BaseReader): api_token (str): Confluence API token, see https://confluence.atlassian.com/cloud/api-tokens-938839638.html user_name (str): Confluence username, used for basic auth. Must be used with `password`. password (str): Confluence password, used for basic auth. Must be used with `user_name`. + client_args (dict): Additional keyword arguments to pass directly to the Atlassian Confluence client constructor, for example `{'backoff_and_retry': True}`. """ @@ -44,6 +45,7 @@ class ConfluenceReader(BaseReader): api_token: Optional[str] = None, user_name: Optional[str] = None, password: Optional[str] = None, + client_args: Optional[dict] = None, ) -> None: if base_url is None: raise ValueError("Must provide `base_url`") @@ -58,20 +60,30 @@ class ConfluenceReader(BaseReader): " atlassian-python-api`" ) self.confluence: Confluence = None + if client_args is None: + client_args = {} if oauth2: - self.confluence = Confluence(url=base_url, oauth2=oauth2, cloud=cloud) + self.confluence = Confluence( + url=base_url, oauth2=oauth2, cloud=cloud, **client_args + ) else: if api_token is not None: - self.confluence = Confluence(url=base_url, token=api_token, cloud=cloud) + self.confluence = Confluence( + url=base_url, token=api_token, cloud=cloud, **client_args + ) elif user_name is not None and password is not None: self.confluence = Confluence( - url=base_url, username=user_name, password=password, cloud=cloud + url=base_url, + username=user_name, + password=password, + cloud=cloud, + **client_args, ) else: api_token = os.getenv(CONFLUENCE_API_TOKEN) if api_token is not None: self.confluence = Confluence( - url=base_url, token=api_token, cloud=cloud + url=base_url, token=api_token, cloud=cloud, **client_args ) else: user_name = os.getenv(CONFLUENCE_USERNAME) @@ -82,6 +94,7 @@ class ConfluenceReader(BaseReader): username=user_name, password=password, cloud=cloud, + **client_args, ) else: raise ValueError( diff --git a/llama-index-integrations/readers/llama-index-readers-confluence/pyproject.toml b/llama-index-integrations/readers/llama-index-readers-confluence/pyproject.toml index 21cb68aed91df7c933b7d8ec12a15ba677135fd9..a85030add9e6a64205928824d1abad4b00f99dde 100644 --- a/llama-index-integrations/readers/llama-index-readers-confluence/pyproject.toml +++ b/llama-index-integrations/readers/llama-index-readers-confluence/pyproject.toml @@ -28,7 +28,7 @@ license = "MIT" maintainers = ["zywilliamli"] name = "llama-index-readers-confluence" readme = "README.md" -version = "0.2.1" +version = "0.2.2" [tool.poetry.dependencies] python = ">=3.8.1,<4.0" diff --git a/llama-index-integrations/readers/llama-index-readers-confluence/tests/test_readers_confluence.py b/llama-index-integrations/readers/llama-index-readers-confluence/tests/test_readers_confluence.py index 6a8d9d2b1e23bcd835282b7bc9b5d23287a48c91..8040a37515970533c36950689a45b74586c720e3 100644 --- a/llama-index-integrations/readers/llama-index-readers-confluence/tests/test_readers_confluence.py +++ b/llama-index-integrations/readers/llama-index-readers-confluence/tests/test_readers_confluence.py @@ -1,3 +1,5 @@ +from unittest.mock import patch + import pytest from llama_index.readers.confluence import ConfluenceReader @@ -31,6 +33,22 @@ def test_confluence_reader_with_api_token(): assert reader.confluence is not None +def test_confluence_reader_with_client_args(): + with patch("atlassian.Confluence") as MockConstructor: + reader = ConfluenceReader( + base_url="https://example.atlassian.net/wiki", + api_token="example_api_token", + client_args={"backoff_and_retry": True}, + ) + assert reader.confluence is not None + MockConstructor.assert_called_once_with( + url="https://example.atlassian.net/wiki", + token="example_api_token", + cloud=True, + backoff_and_retry=True, + ) + + def test_confluence_reader_with_basic_auth(): reader = ConfluenceReader( base_url="https://example.atlassian.net/wiki",