Support Confluence cookies

langchain-ai · Dec 17, 2024 · 403c4ee · 403c4ee
1 parent cdf6202
commit 403c4ee
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 5 deletions.
diff --git a/docs/docs/integrations/document_loaders/confluence.ipynb b/docs/docs/integrations/document_loaders/confluence.ipynb
@@ -11,7 +11,7 @@
     "A loader for `Confluence` pages.\n",
     "\n",
     "\n",
-    "This currently supports `username/api_key`, `Oauth2 login`. Additionally, on-prem installations also support `token` authentication. \n",
+    "This currently supports `username/api_key`, `Oauth2 login`, `cookies`. Additionally, on-prem installations also support `token` authentication. \n",
     "\n",
     "\n",
     "Specify a list `page_id`-s and/or `space_key` to load in the corresponding pages into Document objects, if both are specified the union of both sets will be returned.\n",

diff --git a/libs/community/langchain_community/document_loaders/confluence.py b/libs/community/langchain_community/document_loaders/confluence.py
@@ -34,8 +34,8 @@ class ConfluenceLoader(BaseLoader):
     """Load `Confluence` pages.
 
     Port of https://llamahub.ai/l/confluence
-    This currently supports username/api_key, Oauth2 login or personal access token
-    authentication.
+    This currently supports username/api_key, Oauth2 login, personal access token
+    or cookies authentication.
 
     Specify a list page_ids and/or space_key to load in the corresponding pages into
     Document objects, if both are specified the union of both sets will be returned.
@@ -93,6 +93,8 @@ class ConfluenceLoader(BaseLoader):
     :type oauth2: dict, optional
     :param token: _description_, defaults to None
     :type token: str, optional
+    :param cookies: _description_, defaults to {}
+    :type cookies: dict, optional
     :param cloud: _description_, defaults to True
     :type cloud: bool, optional
     :param number_of_retries: How many times to retry, defaults to 3
@@ -152,6 +154,7 @@ def __init__(
         session: Optional[requests.Session] = None,
         oauth2: Optional[dict] = None,
         token: Optional[str] = None,
+        cookies: Optional[dict] = None,
         cloud: Optional[bool] = True,
         number_of_retries: Optional[int] = 3,
         min_retry_seconds: Optional[int] = 2,
@@ -197,6 +200,7 @@ def __init__(
             username=username,
             session=session,
             oauth2=oauth2,
+            cookies=cookies,
             token=token,
         )
         if errors:
@@ -224,6 +228,10 @@ def __init__(
             self.confluence = Confluence(
                 url=url, token=token, cloud=cloud, **confluence_kwargs
             )
+        elif cookies:
+            self.confluence = Confluence(
+                url=url, cookies=cookies, cloud=cloud, **confluence_kwargs
+            )
         else:
             self.confluence = Confluence(
                 url=url,
@@ -241,6 +249,7 @@ def validate_init_args(
         session: Optional[requests.Session] = None,
         oauth2: Optional[dict] = None,
         token: Optional[str] = None,
+        cookies: Optional[dict] = None,
     ) -> Union[List, None]:
         """Validates proper combinations of init arguments"""
 
@@ -255,10 +264,10 @@ def validate_init_args(
             )
 
         non_null_creds = list(
-            x is not None for x in ((api_key or username), session, oauth2, token)
+            x is not None for x in ((api_key or username), session, oauth2, token, cookies)
         )
         if sum(non_null_creds) > 1:
-            all_names = ("(api_key, username)", "session", "oauth2", "token")
+            all_names = ("(api_key, username)", "session", "oauth2", "token", "cookies")
             provided = tuple(n for x, n in zip(non_null_creds, all_names) if x)
             errors.append(
                 f"Cannot provide a value for more than one of: {all_names}. Received "

diff --git a/libs/community/tests/unit_tests/document_loaders/test_confluence.py b/libs/community/tests/unit_tests/document_loaders/test_confluence.py
@@ -68,6 +68,16 @@ def test_confluence_loader_initialization_invalid(self) -> None:
                 session=requests.Session(),
             )
 
+        with pytest.raises(ValueError):
+            ConfluenceLoader(
+                self.CONFLUENCE_URL,
+                username=self.MOCK_USERNAME,
+                api_key=self.MOCK_API_TOKEN,
+                cookies={
+                    "key": "value",
+                },
+            )
+
     def test_confluence_loader_initialization_from_env(
         self, mock_confluence: MagicMock
     ) -> None: