From 403c4eedb4de68ec9c5563b3b2364139f4af891d Mon Sep 17 00:00:00 2001
From: soulmachine <soulmachine@gmail.com>
Date: Tue, 17 Dec 2024 00:39:16 -0800
Subject: [PATCH] Support Confluence cookies

---
 .../document_loaders/confluence.ipynb           |  2 +-
 .../document_loaders/confluence.py              | 17 +++++++++++++----
 .../document_loaders/test_confluence.py         | 10 ++++++++++
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/docs/docs/integrations/document_loaders/confluence.ipynb b/docs/docs/integrations/document_loaders/confluence.ipynb
index c7f04f0bae558..d325e129d08d7 100644
--- 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
index 954f4139cae5c..d57b8ba58e44f 100644
--- 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
index abb47326beef7..8075abdee1939 100644
--- 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: